diff --git a/c/CMakeLists.txt b/c/CMakeLists.txt index 59932c32ee..37bad92c3b 100644 --- a/c/CMakeLists.txt +++ b/c/CMakeLists.txt @@ -86,6 +86,7 @@ add_library( cuvs_c SHARED src/core/c_api.cpp src/cluster/kmeans.cpp + src/cluster/soar.cpp $<$:src/cluster/mg_kmeans.cpp> src/neighbors/brute_force.cpp src/neighbors/ivf_flat.cpp diff --git a/c/include/cuvs/cluster/soar.h b/c/include/cuvs/cluster/soar.h new file mode 100644 index 0000000000..3ffb76b5d4 --- /dev/null +++ b/c/include/cuvs/cluster/soar.h @@ -0,0 +1,131 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include + +#include + +#ifdef __cplusplus +extern "C" { +#endif + +/** + * @defgroup soar_c_params SOAR hyperparameters + * @{ + */ + +/** + * @brief Hyper-parameters for SOAR assignment. + */ +struct cuvsSoarParams { + /** + * Weight of the projection of the secondary residual onto the primary residual in the SOAR + * loss. Larger values penalize secondary centroids whose residual is aligned with the primary + * residual, favoring complementary assignments. `0` reduces the loss to plain squared distance, + * which the primary centroid itself minimizes, so nothing is spilled. Default: 1.0. + */ + float lambda; +}; + +typedef struct cuvsSoarParams* cuvsSoarParams_t; + +/** + * @brief Allocate SOAR params, and populate with default values + * + * @param[out] params cuvsSoarParams_t to allocate + * @return cuvsError_t + */ +CUVS_EXPORT cuvsError_t cuvsSoarParamsCreate(cuvsSoarParams_t* params); + +/** + * @brief De-allocate SOAR params + * + * @param[in] params cuvsSoarParams_t to de-allocate + * @return cuvsError_t + */ +CUVS_EXPORT cuvsError_t cuvsSoarParamsDestroy(cuvsSoarParams_t params); + +/** + * @} + */ + +/** + * @defgroup soar_c SOAR assignment + * @{ + */ + +/** + * @brief Assign a secondary ("spilled") cluster to each row of the dataset. + * + * SOAR (Spilling with Orthogonality-Amplified Residuals) picks, for each vector, a second + * centroid that complements the primary assignment instead of merely being the next-closest + * one. It minimizes the loss of Theorem 3.1 of https://arxiv.org/abs/2404.00774: for a vector + * `x` with primary residual `r = x - centroids[labels[i]]`, + * + * `score(c) = ||x - c||^2 + lambda * (dot(r / ||r||, x - c))^2` + * + * and `soar_labels[i]` is the centroid minimizing that score. Indexing a vector under both its + * primary and its secondary centroid improves recall for queries near a partition boundary. + * + * All tensors must be on device memory. `dataset` and `centroids` must be row-major float32. + * `labels` and `soar_labels` must have the same dtype, either uint32 or int32; int32 is accepted + * so that the output of `cuvsKMeansPredict` can be passed through without a conversion. + * + * The primary centroid is not excluded from the search, so `soar_labels[i] == labels[i]` is a + * possible (and meaningful) result: it says that no other centroid is worth spilling to, which + * is the common case for vectors in the interior of a cluster. + * + * Scratch memory scales as `n_rows * n_clusters * 4` bytes because scores against all centroids + * are materialized at once and are not tiled. Process the dataset in row batches to bound the + * peak device memory usage. + * + * @code{.c} + * #include + * #include + * + * cuvsResources_t res; + * cuvsResourcesCreate(&res); + * + * cuvsSoarParams_t params; + * cuvsSoarParamsCreate(¶ms); + * + * // dataset, centroids and labels come from a prior k-means fit and predict + * cuvsSoarPredict(res, params, &dataset, ¢roids, &labels, &soar_labels); + * + * cuvsSoarParamsDestroy(params); + * cuvsResourcesDestroy(res); + * @endcode + * + * @param[in] res opaque C handle + * @param[in] params Parameters for SOAR assignment. + * @param[in] dataset The dataset. The data must be in row-major format. + * [dim = n_rows x n_features] + * @param[in] centroids Cluster centroids. The data must be in row-major format. + * [dim = n_clusters x n_features] + * @param[in] labels Index of the primary cluster each row belongs to, as produced by + * k-means prediction. Every value must be in `[0, n_clusters)`. + * [len = n_rows] + * @param[out] soar_labels Index of the secondary cluster each row is spilled to. + * [len = n_rows] + * @return cuvsError_t + */ +CUVS_EXPORT cuvsError_t cuvsSoarPredict(cuvsResources_t res, + cuvsSoarParams_t params, + DLManagedTensor* dataset, + DLManagedTensor* centroids, + DLManagedTensor* labels, + DLManagedTensor* soar_labels); + +/** + * @} + */ + +#ifdef __cplusplus +} +#endif diff --git a/c/include/cuvs/core/all.h b/c/include/cuvs/core/all.h index 80b95992eb..4621e9be89 100644 --- a/c/include/cuvs/core/all.h +++ b/c/include/cuvs/core/all.h @@ -14,6 +14,7 @@ #include #include +#include #include #include diff --git a/c/src/cluster/soar.cpp b/c/src/cluster/soar.cpp new file mode 100644 index 0000000000..f620595f4d --- /dev/null +++ b/c/src/cluster/soar.cpp @@ -0,0 +1,92 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include + +#include + +#include +#include +#include + +#include "../core/exceptions.hpp" +#include "../core/interop.hpp" + +namespace { + +template +void _predict(cuvsResources_t res, + const cuvsSoarParams& params, + DLManagedTensor* dataset_tensor, + DLManagedTensor* centroids_tensor, + DLManagedTensor* labels_tensor, + DLManagedTensor* soar_labels_tensor) +{ + auto res_ptr = reinterpret_cast(res); + + using matrix_type = raft::device_matrix_view; + + auto dataset = cuvs::core::from_dlpack(dataset_tensor); + auto centroids = cuvs::core::from_dlpack(centroids_tensor); + auto labels = + cuvs::core::from_dlpack>(labels_tensor); + auto soar_labels = + cuvs::core::from_dlpack>(soar_labels_tensor); + + // Cluster ids are non-negative and the C++ API takes uint32, so an int32 label array has the + // same bit pattern and can be handed over without a conversion pass over the labels. + auto labels_u32 = raft::make_device_vector_view( + reinterpret_cast(labels.data_handle()), labels.extent(0)); + auto soar_labels_u32 = raft::make_device_vector_view( + reinterpret_cast(soar_labels.data_handle()), soar_labels.extent(0)); + + cuvs::cluster::soar::params cpp_params; + cpp_params.lambda = params.lambda; + + cuvs::cluster::soar::predict( + *res_ptr, cpp_params, dataset, centroids, labels_u32, soar_labels_u32); +} + +} // namespace + +extern "C" cuvsError_t cuvsSoarParamsCreate(cuvsSoarParams_t* params) +{ + return cuvs::core::translate_exceptions([=] { + cuvs::cluster::soar::params cpp_params; + *params = new cuvsSoarParams{.lambda = cpp_params.lambda}; + }); +} + +extern "C" cuvsError_t cuvsSoarParamsDestroy(cuvsSoarParams_t params) +{ + return cuvs::core::translate_exceptions([=] { delete params; }); +} + +extern "C" cuvsError_t cuvsSoarPredict(cuvsResources_t res, + cuvsSoarParams_t params, + DLManagedTensor* dataset, + DLManagedTensor* centroids, + DLManagedTensor* labels, + DLManagedTensor* soar_labels) +{ + return cuvs::core::translate_exceptions([=] { + auto labels_dtype = labels->dl_tensor.dtype; + auto soar_labels_dtype = soar_labels->dl_tensor.dtype; + + RAFT_EXPECTS(labels_dtype.code == soar_labels_dtype.code && + labels_dtype.bits == soar_labels_dtype.bits, + "labels and soar_labels must have the same dtype"); + + if (labels_dtype.code == kDLUInt && labels_dtype.bits == 32) { + _predict(res, *params, dataset, centroids, labels, soar_labels); + } else if (labels_dtype.code == kDLInt && labels_dtype.bits == 32) { + _predict(res, *params, dataset, centroids, labels, soar_labels); + } else { + RAFT_FAIL("Unsupported labels DLtensor dtype: %d and bits: %d", + labels_dtype.code, + labels_dtype.bits); + } + }); +} diff --git a/c/tests/CMakeLists.txt b/c/tests/CMakeLists.txt index 7d6c588bd9..78ba4816e1 100644 --- a/c/tests/CMakeLists.txt +++ b/c/tests/CMakeLists.txt @@ -80,6 +80,7 @@ ConfigureTest( NAME DISTANCE_C_TEST PATH distance/run_pairwise_distance_c.c distance/pairwise_distance_c.cu ) ConfigureTest(NAME KMEANS_C_TEST PATH cluster/kmeans_c.cu) +ConfigureTest(NAME SOAR_C_TEST PATH cluster/soar_c.cu) ConfigureTest(NAME SELECT_K_C_TEST PATH selection/select_k_c.cu) if(BUILD_MG_ALGOS) ConfigureTest(NAME KMEANS_MG_C_TEST PATH cluster/kmeans_mg_c.cu) diff --git a/c/tests/cluster/soar_c.cu b/c/tests/cluster/soar_c.cu new file mode 100644 index 0000000000..4101f7648a --- /dev/null +++ b/c/tests/cluster/soar_c.cu @@ -0,0 +1,330 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include + +#include +#include +#include +#include +#include + +#include "../../src/core/interop.hpp" +#include + +#include +#include +#include +#include +#include +#include + +namespace { + +constexpr int64_t kNRows = 512; +constexpr int64_t kDim = 16; +constexpr int64_t kNClusters = 24; + +template +DLManagedTensor make_matrix_tensor(T* data, int64_t rows, int64_t cols) +{ + DLManagedTensor tensor{}; + cuvs::core::to_dlpack(raft::make_device_matrix_view(data, rows, cols), + &tensor); + return tensor; +} + +template +DLManagedTensor make_vector_tensor(T* data, int64_t size) +{ + DLManagedTensor tensor{}; + cuvs::core::to_dlpack(raft::make_device_vector_view(data, size), &tensor); + return tensor; +} + +void free_tensor(DLManagedTensor& t) +{ + if (t.deleter) { t.deleter(&t); } +} + +/** Uniform random matrix in [-1, 1], generated on the host so the tests are reproducible. */ +std::vector random_matrix(int64_t n_rows, int64_t dim, uint64_t seed) +{ + std::mt19937_64 rng(seed); + std::uniform_real_distribution dist(-1.0f, 1.0f); + std::vector data(n_rows * dim); + std::generate(data.begin(), data.end(), [&]() { return dist(rng); }); + return data; +} + +/** Index of the closest centroid in L2, i.e. what k-means prediction would produce. */ +std::vector nearest_centroids(const std::vector& dataset, + const std::vector& centroids, + int64_t n_rows, + int64_t dim, + int64_t n_clusters) +{ + std::vector labels(n_rows); + for (int64_t i = 0; i < n_rows; i++) { + double best_distance = std::numeric_limits::max(); + for (int64_t c = 0; c < n_clusters; c++) { + double distance = 0.0; + for (int64_t k = 0; k < dim; k++) { + double diff = static_cast(dataset[i * dim + k]) - centroids[c * dim + k]; + distance += diff * diff; + } + if (distance < best_distance) { + best_distance = distance; + labels[i] = static_cast(c); + } + } + } + return labels; +} + +/** + * `||x - c||^2 + lambda * (dot(r / ||r||, x - c))^2`, the loss the implementation minimizes over + * all centroids, up to a per-row constant that does not move the argmin. + */ +double soar_score(const float* x, + const float* residual, + const float* centroid, + int64_t dim, + float lambda) +{ + double residual_norm = 0.0; + for (int64_t k = 0; k < dim; k++) { + residual_norm += static_cast(residual[k]) * residual[k]; + } + residual_norm = std::sqrt(residual_norm); + + double squared_distance = 0.0; + double projection = 0.0; + for (int64_t k = 0; k < dim; k++) { + double diff = static_cast(x[k]) - centroid[k]; + squared_distance += diff * diff; + projection += diff * (residual[k] / residual_norm); + } + + return squared_distance + static_cast(lambda) * projection * projection; +} + +/** + * Holds the random dataset, the primary k-means labels and the matching device buffers shared by + * the tests below. + */ +struct SoarFixture { + explicit SoarFixture(rmm::cuda_stream_view stream) + : dataset(kNRows * kDim, stream), + centroids(kNClusters * kDim, stream), + labels(kNRows, stream), + soar_labels(kNRows, stream) + { + h_dataset = random_matrix(kNRows, kDim, 1234ULL); + h_centroids = random_matrix(kNClusters, kDim, 5678ULL); + h_labels = nearest_centroids(h_dataset, h_centroids, kNRows, kDim, kNClusters); + + raft::update_device(dataset.data(), h_dataset.data(), h_dataset.size(), stream); + raft::update_device(centroids.data(), h_centroids.data(), h_centroids.size(), stream); + raft::update_device(labels.data(), h_labels.data(), h_labels.size(), stream); + } + + std::vector h_dataset; + std::vector h_centroids; + std::vector h_labels; + + rmm::device_uvector dataset; + rmm::device_uvector centroids; + rmm::device_uvector labels; + rmm::device_uvector soar_labels; +}; + +} // namespace + +TEST(SoarC, PredictMatchesExhaustiveHostSearch) +{ + raft::device_resources handle; + auto stream = raft::resource::get_cuda_stream(handle); + + SoarFixture fixture(stream); + handle.sync_stream(); + + cuvsResources_t res; + ASSERT_EQ(cuvsResourcesCreate(&res), CUVS_SUCCESS); + + cuvsSoarParams_t params; + ASSERT_EQ(cuvsSoarParamsCreate(¶ms), CUVS_SUCCESS); + ASSERT_FLOAT_EQ(params->lambda, 1.0f) << "default lambda should match the C++ params"; + + // Deliberately not the default, so the test fails if the wrapper drops lambda on its way to + // the C++ implementation. + params->lambda = 2.0f; + + auto dataset_t = make_matrix_tensor(fixture.dataset.data(), kNRows, kDim); + auto centroids_t = make_matrix_tensor(fixture.centroids.data(), kNClusters, kDim); + auto labels_t = make_vector_tensor(fixture.labels.data(), kNRows); + auto soar_labels_t = make_vector_tensor(fixture.soar_labels.data(), kNRows); + + ASSERT_EQ(cuvsSoarPredict(res, params, &dataset_t, ¢roids_t, &labels_t, &soar_labels_t), + CUVS_SUCCESS) + << cuvsGetLastErrorText(); + + ASSERT_EQ(cuvsStreamSync(res), CUVS_SUCCESS); + + std::vector h_soar_labels(kNRows); + raft::update_host(h_soar_labels.data(), fixture.soar_labels.data(), kNRows, stream); + handle.sync_stream(); + + // Compare losses rather than ids so the test is not fragile when two centroids tie. + std::vector residual(kDim); + int64_t n_spilled = 0; + for (int64_t i = 0; i < kNRows; i++) { + ASSERT_LT(h_soar_labels[i], kNClusters) << "row " << i << " got an out-of-range label"; + if (h_soar_labels[i] != fixture.h_labels[i]) { ++n_spilled; } + + const float* x = &fixture.h_dataset[i * kDim]; + const float* primary = &fixture.h_centroids[fixture.h_labels[i] * kDim]; + for (int64_t k = 0; k < kDim; k++) { + residual[k] = x[k] - primary[k]; + } + + double best_score = std::numeric_limits::max(); + for (int64_t c = 0; c < kNClusters; c++) { + best_score = std::min( + best_score, + soar_score(x, residual.data(), &fixture.h_centroids[c * kDim], kDim, params->lambda)); + } + + double actual_score = soar_score( + x, residual.data(), &fixture.h_centroids[h_soar_labels[i] * kDim], kDim, params->lambda); + + ASSERT_NEAR(actual_score, best_score, 1e-3 * std::max(1.0, std::abs(best_score))) + << "row " << i << " was not assigned a loss-minimizing centroid"; + } + + // Guards against the loss check above going vacuous: if the fixture ever stopped spilling, + // simply echoing the primary labels would satisfy it. + EXPECT_GT(n_spilled, 0) << "expected at least one row to spill"; + + free_tensor(dataset_t); + free_tensor(centroids_t); + free_tensor(labels_t); + free_tensor(soar_labels_t); + ASSERT_EQ(cuvsSoarParamsDestroy(params), CUVS_SUCCESS); + ASSERT_EQ(cuvsResourcesDestroy(res), CUVS_SUCCESS); +} + +TEST(SoarC, Int32LabelsMatchUint32Labels) +{ + raft::device_resources handle; + auto stream = raft::resource::get_cuda_stream(handle); + + SoarFixture fixture(stream); + + // The same primary labels, as int32, which is what cuvsKMeansPredict writes. + std::vector h_labels_i32(fixture.h_labels.begin(), fixture.h_labels.end()); + rmm::device_uvector labels_i32(kNRows, stream); + rmm::device_uvector soar_labels_i32(kNRows, stream); + raft::update_device(labels_i32.data(), h_labels_i32.data(), h_labels_i32.size(), stream); + handle.sync_stream(); + + cuvsResources_t res; + ASSERT_EQ(cuvsResourcesCreate(&res), CUVS_SUCCESS); + cuvsSoarParams_t params; + ASSERT_EQ(cuvsSoarParamsCreate(¶ms), CUVS_SUCCESS); + + auto dataset_t = make_matrix_tensor(fixture.dataset.data(), kNRows, kDim); + auto centroids_t = make_matrix_tensor(fixture.centroids.data(), kNClusters, kDim); + + auto labels_u32_t = make_vector_tensor(fixture.labels.data(), kNRows); + auto soar_labels_u32_t = make_vector_tensor(fixture.soar_labels.data(), kNRows); + ASSERT_EQ( + cuvsSoarPredict(res, params, &dataset_t, ¢roids_t, &labels_u32_t, &soar_labels_u32_t), + CUVS_SUCCESS) + << cuvsGetLastErrorText(); + + auto labels_i32_t = make_vector_tensor(labels_i32.data(), kNRows); + auto soar_labels_i32_t = make_vector_tensor(soar_labels_i32.data(), kNRows); + ASSERT_EQ( + cuvsSoarPredict(res, params, &dataset_t, ¢roids_t, &labels_i32_t, &soar_labels_i32_t), + CUVS_SUCCESS) + << cuvsGetLastErrorText(); + + ASSERT_EQ(cuvsStreamSync(res), CUVS_SUCCESS); + + std::vector h_from_u32(kNRows); + std::vector h_from_i32(kNRows); + raft::update_host(h_from_u32.data(), fixture.soar_labels.data(), kNRows, stream); + raft::update_host(h_from_i32.data(), soar_labels_i32.data(), kNRows, stream); + handle.sync_stream(); + + for (int64_t i = 0; i < kNRows; i++) { + ASSERT_EQ(static_cast(h_from_u32[i]), h_from_i32[i]) + << "row " << i << " differs between the uint32 and int32 label paths"; + } + + free_tensor(dataset_t); + free_tensor(centroids_t); + free_tensor(labels_u32_t); + free_tensor(soar_labels_u32_t); + free_tensor(labels_i32_t); + free_tensor(soar_labels_i32_t); + ASSERT_EQ(cuvsSoarParamsDestroy(params), CUVS_SUCCESS); + ASSERT_EQ(cuvsResourcesDestroy(res), CUVS_SUCCESS); +} + +TEST(SoarC, RejectsInvalidInputs) +{ + raft::device_resources handle; + auto stream = raft::resource::get_cuda_stream(handle); + + SoarFixture fixture(stream); + rmm::device_uvector soar_labels_i32(kNRows, stream); + rmm::device_uvector labels_f64(kNRows, stream); + handle.sync_stream(); + + cuvsResources_t res; + ASSERT_EQ(cuvsResourcesCreate(&res), CUVS_SUCCESS); + cuvsSoarParams_t params; + ASSERT_EQ(cuvsSoarParamsCreate(¶ms), CUVS_SUCCESS); + + auto dataset_t = make_matrix_tensor(fixture.dataset.data(), kNRows, kDim); + auto centroids_t = make_matrix_tensor(fixture.centroids.data(), kNClusters, kDim); + auto labels_t = make_vector_tensor(fixture.labels.data(), kNRows); + auto soar_labels_t = make_vector_tensor(fixture.soar_labels.data(), kNRows); + + // labels and soar_labels must agree on dtype. + auto soar_labels_i32_t = make_vector_tensor(soar_labels_i32.data(), kNRows); + EXPECT_EQ(cuvsSoarPredict(res, params, &dataset_t, ¢roids_t, &labels_t, &soar_labels_i32_t), + CUVS_ERROR); + + // float64 is not a valid label dtype. + auto labels_f64_t = make_vector_tensor(labels_f64.data(), kNRows); + EXPECT_EQ(cuvsSoarPredict(res, params, &dataset_t, ¢roids_t, &labels_f64_t, &labels_f64_t), + CUVS_ERROR); + + // The centroid dimensionality has to match the dataset. + auto narrow_centroids_t = make_matrix_tensor(fixture.centroids.data(), kNClusters, kDim - 1); + EXPECT_EQ( + cuvsSoarPredict(res, params, &dataset_t, &narrow_centroids_t, &labels_t, &soar_labels_t), + CUVS_ERROR); + + // soar_labels has to be as long as the dataset. + auto short_soar_labels_t = make_vector_tensor(fixture.soar_labels.data(), kNRows - 1); + EXPECT_EQ( + cuvsSoarPredict(res, params, &dataset_t, ¢roids_t, &labels_t, &short_soar_labels_t), + CUVS_ERROR); + + free_tensor(dataset_t); + free_tensor(centroids_t); + free_tensor(labels_t); + free_tensor(soar_labels_t); + free_tensor(soar_labels_i32_t); + free_tensor(labels_f64_t); + free_tensor(narrow_centroids_t); + free_tensor(short_soar_labels_t); + ASSERT_EQ(cuvsSoarParamsDestroy(params), CUVS_SUCCESS); + ASSERT_EQ(cuvsResourcesDestroy(res), CUVS_SUCCESS); +} diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index ef9977f72c..6757743f9e 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -1432,6 +1432,7 @@ if(NOT BUILD_CPU_ONLY) src/cluster/kmeans_transform_double.cu src/cluster/kmeans_transform_float.cu src/cluster/single_linkage_float.cu + src/cluster/soar.cu src/cluster/spectral.cu src/core/bitset.cu src/core/bloom_filter.cu diff --git a/cpp/include/cuvs/cluster/soar.hpp b/cpp/include/cuvs/cluster/soar.hpp new file mode 100644 index 0000000000..4c082e9ce4 --- /dev/null +++ b/cpp/include/cuvs/cluster/soar.hpp @@ -0,0 +1,132 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include + +#include + +namespace CUVS_EXPORT cuvs { +namespace cluster { +namespace soar { + +/** + * @defgroup soar_params SOAR hyperparameters + * @{ + */ + +/** + * Simple object to specify hyper-parameters for SOAR assignment. + */ +struct params { + /** + * Weight of the projection of the secondary residual onto the primary residual in the SOAR + * loss. Larger values penalize secondary centroids whose residual is aligned with the primary + * residual, favoring complementary assignments. `0` reduces the loss to plain squared distance, + * which the primary centroid itself minimizes, so nothing is spilled. Default: 1.0. + */ + float lambda = 1.0f; +}; + +/** + * @} + */ + +/** + * @defgroup soar_predict SOAR assignment + * @{ + */ + +/** + * @brief Assign a secondary ("spilled") cluster to each row of the dataset. + * + * SOAR (Spilling with Orthogonality-Amplified Residuals) picks, for each vector, a second + * centroid that complements the primary assignment instead of merely being the next-closest + * one. It minimizes the loss of Theorem 3.1 of https://arxiv.org/abs/2404.00774: for a vector + * `x` with primary residual `r = x - centroids[labels[i]]`, + * + * `score(c) = ||x - c||^2 + lambda * (dot(r / ||r||, x - c))^2` + * + * and `soar_labels[i]` is the centroid minimizing that score. Indexing a vector under both its + * primary and its secondary centroid improves recall for queries near a partition boundary. + * + * Only float32 data and uint32 labels are supported. + * + * The primary centroid is not excluded from the search, so `soar_labels[i] == labels[i]` is a + * possible (and meaningful) result: it says that no other centroid is worth spilling to, which + * is the common case for vectors in the interior of a cluster. Callers that treat SOAR as a + * strictly second posting list should test for this case and skip those rows. + * + * Scratch memory scales as `n_rows * n_clusters * 4` bytes because scores against all centroids + * are materialized at once and are not tiled. Process the dataset in row batches to bound the + * peak device memory usage. + * + * @code{.cpp} + * #include + * #include + * #include + * using namespace cuvs::cluster; + * ... + * raft::resources handle; + * cuvs::cluster::kmeans::balanced_params kmeans_params; + * int64_t n_features = 15, n_clusters = 100; + * auto centroids = raft::make_device_matrix(handle, n_clusters, n_features); + * + * // primary assignments, e.g. from balanced k-means + * kmeans::fit(handle, + * kmeans_params, + * dataset, + * centroids.view()); + * ... + * auto labels = raft::make_device_vector(handle, dataset.extent(0)); + * + * kmeans::predict(handle, + * kmeans_params, + * dataset, + * raft::make_const_mdspan(centroids.view()), + * labels.view()); + * ... + * // secondary assignments + * cuvs::cluster::soar::params soar_params; + * auto soar_labels = raft::make_device_vector(handle, dataset.extent(0)); + * + * soar::predict(handle, + * soar_params, + * dataset, + * raft::make_const_mdspan(centroids.view()), + * raft::make_const_mdspan(labels.view()), + * soar_labels.view()); + * // soar_labels now holds one secondary centroid id per row + * @endcode + * + * @param[in] handle The raft handle. + * @param[in] params Parameters for SOAR assignment. + * @param[in] dataset The dataset. The data must be in row-major format. + * [dim = n_rows x n_features] + * @param[in] centroids Cluster centroids. The data must be in row-major format. + * [dim = n_clusters x n_features] + * @param[in] labels Index of the primary cluster each row belongs to, as produced by + * k-means prediction. Every value must be in `[0, n_clusters)`. + * [len = n_rows] + * @param[out] soar_labels Index of the secondary cluster each row is spilled to. + * [len = n_rows] + */ +void predict(raft::resources const& handle, + const soar::params& params, + raft::device_matrix_view dataset, + raft::device_matrix_view centroids, + raft::device_vector_view labels, + raft::device_vector_view soar_labels); + +/** + * @} + */ + +} // namespace soar +} // namespace cluster +} // namespace CUVS_EXPORT cuvs diff --git a/cpp/src/neighbors/scann/detail/scann_soar.cuh b/cpp/src/cluster/detail/soar.cuh similarity index 68% rename from cpp/src/neighbors/scann/detail/scann_soar.cuh rename to cpp/src/cluster/detail/soar.cuh index 38bb7b0858..b039dd3e33 100644 --- a/cpp/src/neighbors/scann/detail/scann_soar.cuh +++ b/cpp/src/cluster/detail/soar.cuh @@ -1,13 +1,16 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ +#pragma once + #include #include #include #include #include +#include #include #include #include @@ -17,27 +20,72 @@ #include #include +namespace cuvs::cluster::soar::detail { + +/** + * @brief Subtract cluster center coordinates from each dataset vector. + * + * residual[i, k] = dataset[i ,k] - centers[l, k], + * where l = labels[i], the cluster label corresponding to vector i. + * + * An identical copy lives in `neighbors/scann/detail/scann_quantize.cuh`, where it is also used + * outside the SOAR path to build the PQ trainset residuals. + * + * @tparam T + * @tparam LabelT + * @param res raft resources + * @param dataset dataset vectors, size [n_rows, dim] + * @param centers cluster center coordinates, size [n_clusters, dim] + * @param labels cluster labels, size [n_rows] + * @return device matrix with the residuals, size [n_rows, dim] + */ +template +auto compute_residuals(raft::resources const& res, + raft::device_matrix_view dataset, + raft::device_matrix_view centers, + raft::device_vector_view labels) + -> raft::device_matrix +{ + auto dim = dataset.extent(1); + auto residuals = raft::make_device_matrix(res, labels.extent(0), dim); + + raft::linalg::map_offset( + res, residuals.view(), [dataset, centers, labels, dim] __device__(size_t i) { + int row_idx = i / dim; + int el_idx = i % dim; + return dataset(row_idx, el_idx) - centers(labels(row_idx), el_idx); + }); + + return residuals; +} + /** * @brief Compute SOAR labels for each dataset vector * * Compute a second, spilled cluster for each dataset vector by minimizing * the loss function in Theorem 3.1 of https://arxiv.org/abs/2404.00774 * + * Residuals are an input (`r = x - centers[labels[i]]`) rather than derived here, so a + * caller that already has them can avoid a second pass over the dataset. + * + * The scratch score matrix is [n_rows, n_clusters] floats and is not tiled, so callers are + * responsible for batching rows. + * * @tparam T - * @tparam LavelT - * @param res raft resources + * @tparam LabelT + * @param dev_resources raft resources * @param dataset the dataset, size [n_rows, dim] * @param residuals the residual vectors r, size [n_rows, dim] * @param centers the cluster centers, size [n_clusters, dim] * @param labels the cluster assignments, size [n_rows] - * @param soar_labels the computed soar labels + * @param soar_labels the computed soar labels, size [n_rows] * @param lambda the weight for the projection of a residual r' onto r in the SOAR loss */ template void compute_soar_labels(raft::resources const& dev_resources, raft::device_matrix_view dataset, raft::device_matrix_view residuals, - raft::device_matrix_view centers, + raft::device_matrix_view centers, raft::device_vector_view labels, raft::device_vector_view soar_labels, float lambda) @@ -47,7 +95,6 @@ void compute_soar_labels(raft::resources const& dev_resources, // compute SOAR metric for each center auto soar_scores = raft::make_device_matrix(dev_resources, dataset.extent(0), centers.extent(0)); - auto n_centers = centers.extent(0); auto residuals_norm = raft::make_device_matrix( dev_resources, residuals.extent(0), residuals.extent(1)); @@ -90,15 +137,15 @@ void compute_soar_labels(raft::resources const& dev_resources, auto centers_transpose = raft::make_device_matrix(dev_resources, centers.extent(1), centers.extent(0)); - raft::linalg::reduce(dev_resources, - raft::make_const_mdspan(centers), - centers_norm.view(), - 0.0f, - false, - raft::sq_op(), - raft::add_op()); + raft::linalg::reduce( + dev_resources, centers, centers_norm.view(), 0.0f, false, raft::sq_op(), raft::add_op()); - raft::linalg::transpose(dev_resources, centers, centers_transpose.view()); + // raft::linalg::transpose requires input and output views of the same type; it does not + // write to the input. + auto nc_centers = raft::make_device_matrix_view( + const_cast(centers.data_handle()), centers.extent(0), centers.extent(1)); + + raft::linalg::transpose(dev_resources, nc_centers, centers_transpose.view()); raft::linalg::gemm( dev_resources, residuals_norm.view(), centers_transpose.view(), soar_scores.view()); @@ -146,3 +193,5 @@ void compute_soar_labels(raft::resources const& dev_resources, raft::matrix::argmin(dev_resources, raft::make_const_mdspan(soar_scores.view()), soar_labels); } + +} // namespace cuvs::cluster::soar::detail diff --git a/cpp/src/cluster/soar.cu b/cpp/src/cluster/soar.cu new file mode 100644 index 0000000000..804755003f --- /dev/null +++ b/cpp/src/cluster/soar.cu @@ -0,0 +1,54 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "detail/soar.cuh" + +#include + +#include + +namespace cuvs::cluster::soar { + +void predict(raft::resources const& handle, + const soar::params& params, + raft::device_matrix_view dataset, + raft::device_matrix_view centroids, + raft::device_vector_view labels, + raft::device_vector_view soar_labels) +{ + int64_t n_rows = dataset.extent(0); + int64_t dim = dataset.extent(1); + int64_t n_clusters = centroids.extent(0); + + RAFT_EXPECTS(centroids.extent(1) == dim, + "Number of features in the dataset (%zd) and in the centroids (%zd) must match.", + dim, + centroids.extent(1)); + RAFT_EXPECTS(n_clusters > 0, "The number of centroids must be positive."); + RAFT_EXPECTS(dim > 0, "The number of features must be positive."); + RAFT_EXPECTS(labels.extent(0) == n_rows, + "The number of labels (%zd) must match the number of rows in the dataset (%zd).", + labels.extent(0), + n_rows); + RAFT_EXPECTS( + soar_labels.extent(0) == n_rows, + "The number of soar labels (%zd) must match the number of rows in the dataset (%zd).", + soar_labels.extent(0), + n_rows); + + if (n_rows == 0) { return; } + + auto residuals = detail::compute_residuals(handle, dataset, centroids, labels); + + detail::compute_soar_labels(handle, + dataset, + raft::make_const_mdspan(residuals.view()), + centroids, + labels, + soar_labels, + params.lambda); +} + +} // namespace cuvs::cluster::soar diff --git a/cpp/src/neighbors/scann/detail/scann_build.cuh b/cpp/src/neighbors/scann/detail/scann_build.cuh index f4e2234c75..4af469c0a0 100644 --- a/cpp/src/neighbors/scann/detail/scann_build.cuh +++ b/cpp/src/neighbors/scann/detail/scann_build.cuh @@ -24,10 +24,10 @@ #include #include +#include "../../../cluster/detail/soar.cuh" #include "scann_avq.cuh" #include "scann_common.cuh" #include "scann_quantize.cuh" -#include "scann_soar.cuh" namespace cuvs::neighbors::experimental::scann::detail { using namespace cuvs::spatial::knn::detail; // NOLINT @@ -197,13 +197,14 @@ index build( // Compute SOAR labels. // We compute SOAR labels in this loop to eliminate one HtoD copy of the full dataset. - compute_soar_labels(res, - batch_view, - raft::make_const_mdspan(avq_residuals.view()), - centroids_view, - batch_labels_view, - batch_soar_labels_view, - params.soar_lambda); + cuvs::cluster::soar::detail::compute_soar_labels( + res, + batch_view, + raft::make_const_mdspan(avq_residuals.view()), + raft::make_const_mdspan(centroids_view), + batch_labels_view, + batch_soar_labels_view, + params.soar_lambda); // Compute and quantize residuals using the public PQ API int64_t codes_dim = cuvs::preprocessing::quantize::pq::get_quantized_dim(pq_build_params); diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index d58b13d40f..705dd86e04 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -133,6 +133,7 @@ ConfigureTest( cluster/kmeans_predict_batching.cu cluster/linkage.cu cluster/connect_knn.cu + cluster/soar.cu cluster/spectral.cu GPUS 1 PERCENT 100 diff --git a/cpp/tests/cluster/soar.cu b/cpp/tests/cluster/soar.cu new file mode 100644 index 0000000000..2a5598e32a --- /dev/null +++ b/cpp/tests/cluster/soar.cu @@ -0,0 +1,368 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "../../src/cluster/detail/soar.cuh" +#include "../test_utils.cuh" + +#include + +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include + +namespace cuvs::cluster::soar { + +struct SoarInputs { + int64_t n_rows; + int64_t dim; + int64_t n_clusters; + float lambda; +}; + +::std::ostream& operator<<(::std::ostream& os, const SoarInputs& p) +{ + os << "{ " << p.n_rows << ", " << p.dim << ", " << p.n_clusters << ", " << p.lambda << '}'; + return os; +} + +namespace { + +/** Uniform random matrix in [-1, 1], generated on the host so the tests are reproducible. */ +auto random_matrix(int64_t n_rows, int64_t dim, uint64_t seed) -> std::vector +{ + std::mt19937_64 rng(seed); + std::uniform_real_distribution dist(-1.0f, 1.0f); + std::vector data(n_rows * dim); + std::generate(data.begin(), data.end(), [&]() { return dist(rng); }); + return data; +} + +/** Index of the closest centroid in L2, i.e. what k-means prediction would produce. */ +auto nearest_centroids(const std::vector& dataset, + const std::vector& centroids, + int64_t n_rows, + int64_t dim, + int64_t n_clusters) -> std::vector +{ + std::vector labels(n_rows); + for (int64_t i = 0; i < n_rows; i++) { + double best_distance = std::numeric_limits::max(); + for (int64_t c = 0; c < n_clusters; c++) { + double distance = 0.0; + for (int64_t k = 0; k < dim; k++) { + double diff = static_cast(dataset[i * dim + k]) - centroids[c * dim + k]; + distance += diff * diff; + } + if (distance < best_distance) { + best_distance = distance; + labels[i] = static_cast(c); + } + } + } + return labels; +} + +auto residuals_host(const std::vector& dataset, + const std::vector& centroids, + const std::vector& labels, + int64_t n_rows, + int64_t dim) -> std::vector +{ + std::vector residuals(n_rows * dim); + for (int64_t i = 0; i < n_rows; i++) { + for (int64_t k = 0; k < dim; k++) { + residuals[i * dim + k] = dataset[i * dim + k] - centroids[labels[i] * dim + k]; + } + } + return residuals; +} + +/** + * `||x - c||^2 + lambda * (dot(r / ||r||, x - c))^2`, the loss that the device implementation + * minimizes over all centroids, up to a per-row constant that does not move the argmin. + */ +auto soar_score( + const float* x, const float* residual, const float* centroid, int64_t dim, float lambda) -> double +{ + double residual_norm = 0.0; + for (int64_t k = 0; k < dim; k++) { + residual_norm += static_cast(residual[k]) * residual[k]; + } + residual_norm = std::sqrt(residual_norm); + + double squared_distance = 0.0; + double projection = 0.0; + for (int64_t k = 0; k < dim; k++) { + double diff = static_cast(x[k]) - centroid[k]; + squared_distance += diff * diff; + projection += diff * (residual[k] / residual_norm); + } + + return squared_distance + static_cast(lambda) * projection * projection; +} + +/** The best achievable loss per row, from an exhaustive host search. */ +auto reference_scores(const std::vector& dataset, + const std::vector& centroids, + const std::vector& residuals, + int64_t n_rows, + int64_t dim, + int64_t n_clusters, + float lambda) -> std::vector +{ + std::vector scores(n_rows); + for (int64_t i = 0; i < n_rows; i++) { + double best_score = std::numeric_limits::max(); + for (int64_t c = 0; c < n_clusters; c++) { + best_score = std::min( + best_score, + soar_score(&dataset[i * dim], &residuals[i * dim], ¢roids[c * dim], dim, lambda)); + } + scores[i] = best_score; + } + return scores; +} + +} // namespace + +class SoarTest : public ::testing::TestWithParam { + public: + SoarTest() + : params_(GetParam()), + dataset_(raft::make_device_matrix(handle_, params_.n_rows, params_.dim)), + centroids_( + raft::make_device_matrix(handle_, params_.n_clusters, params_.dim)), + labels_(raft::make_device_vector(handle_, params_.n_rows)), + soar_labels_(raft::make_device_vector(handle_, params_.n_rows)) + { + } + + protected: + void SetUp() override + { + h_dataset_ = random_matrix(params_.n_rows, params_.dim, 1234ULL); + h_centroids_ = random_matrix(params_.n_clusters, params_.dim, 5678ULL); + h_labels_ = + nearest_centroids(h_dataset_, h_centroids_, params_.n_rows, params_.dim, params_.n_clusters); + + auto stream = raft::resource::get_cuda_stream(handle_); + raft::update_device(dataset_.data_handle(), h_dataset_.data(), h_dataset_.size(), stream); + raft::update_device(centroids_.data_handle(), h_centroids_.data(), h_centroids_.size(), stream); + raft::update_device(labels_.data_handle(), h_labels_.data(), h_labels_.size(), stream); + raft::resource::sync_stream(handle_); + } + + /** Run the public API and copy the resulting labels back to the host. */ + auto run_predict() -> std::vector + { + cuvs::cluster::soar::params soar_params; + soar_params.lambda = params_.lambda; + + cuvs::cluster::soar::predict(handle_, + soar_params, + raft::make_const_mdspan(dataset_.view()), + raft::make_const_mdspan(centroids_.view()), + raft::make_const_mdspan(labels_.view()), + soar_labels_.view()); + + return to_host(raft::make_const_mdspan(soar_labels_.view())); + } + + auto to_host(raft::device_vector_view labels) -> std::vector + { + std::vector h_labels(labels.extent(0)); + raft::update_host(h_labels.data(), + labels.data_handle(), + labels.extent(0), + raft::resource::get_cuda_stream(handle_)); + raft::resource::sync_stream(handle_); + return h_labels; + } + + raft::resources handle_; + SoarInputs params_; + + std::vector h_dataset_; + std::vector h_centroids_; + std::vector h_labels_; + + raft::device_matrix dataset_; + raft::device_matrix centroids_; + raft::device_vector labels_; + raft::device_vector soar_labels_; +}; + +/** + * Every label must be a valid centroid id achieving the same loss as an exhaustive host search. + * Comparing losses rather than ids keeps the test from being fragile when two centroids are + * nearly tied. + */ +TEST_P(SoarTest, MatchesHostReference) +{ + auto soar_labels = run_predict(); + + auto h_residuals = + residuals_host(h_dataset_, h_centroids_, h_labels_, params_.n_rows, params_.dim); + auto best_scores = reference_scores(h_dataset_, + h_centroids_, + h_residuals, + params_.n_rows, + params_.dim, + params_.n_clusters, + params_.lambda); + + for (int64_t i = 0; i < params_.n_rows; i++) { + ASSERT_LT(soar_labels[i], static_cast(params_.n_clusters)) + << "label out of range at row " << i; + + double score = soar_score(&h_dataset_[i * params_.dim], + &h_residuals[i * params_.dim], + &h_centroids_[soar_labels[i] * params_.dim], + params_.dim, + params_.lambda); + ASSERT_NEAR(score, best_scores[i], 1e-4 * (1.0 + std::abs(best_scores[i]))) + << "row " << i << " picked centroid " << soar_labels[i]; + } +} + +/** + * The residuals feeding the SOAR loss must match a plain host `x - c[label]`. Both the public + * `predict` and the ScaNN builder depend on this, and `MatchesHostReference` only detects + * residual errors large enough to move an argmin. + */ +TEST_P(SoarTest, ComputeResidualsMatchesHost) +{ + auto h_residuals = + residuals_host(h_dataset_, h_centroids_, h_labels_, params_.n_rows, params_.dim); + + auto residuals = + detail::compute_residuals(handle_, + raft::make_const_mdspan(dataset_.view()), + raft::make_const_mdspan(centroids_.view()), + raft::make_const_mdspan(labels_.view())); + + ASSERT_TRUE(cuvs::devArrMatchHost(h_residuals.data(), + residuals.data_handle(), + h_residuals.size(), + cuvs::CompareApprox(1e-6f), + raft::resource::get_cuda_stream(handle_))); +} + +const std::vector inputs = {{1000, 8, 16, 1.0f}, + {1000, 8, 16, 0.0f}, + {1000, 8, 16, 4.0f}, + {512, 32, 64, 1.5f}, + {17, 3, 2, 1.0f}}; + +INSTANTIATE_TEST_CASE_P(SoarTests, SoarTest, ::testing::ValuesIn(inputs)); + +/** + * A hand-checked case with well-separated centroids, covering both outcomes: a row in the + * interior of its cluster keeps its primary centroid, because no other centroid is close enough + * to be worth spilling to, while a row near a boundary spills to the neighboring cluster. + */ +TEST(SoarTestSmall, SeparatedClusters) +{ + raft::resources handle; + auto stream = raft::resource::get_cuda_stream(handle); + + constexpr int64_t n_rows = 4, dim = 2, n_clusters = 3; + + std::vector h_centroids{0.0f, 0.0f, 100.0f, 0.0f, 0.0f, 100.0f}; + std::vector h_dataset{1.0f, 0.0f, 48.0f, 20.0f, 20.0f, 48.0f, 99.0f, 0.0f}; + std::vector h_labels{0, 0, 0, 1}; + + // Rows 0 and 3 sit next to their own centroid and keep it. Rows 1 and 2 sit between two + // centroids, so the second-closest one wins with a ~13% margin in the loss. + std::vector expected{0, 1, 2, 1}; + + auto dataset = raft::make_device_matrix(handle, n_rows, dim); + auto centroids = raft::make_device_matrix(handle, n_clusters, dim); + auto labels = raft::make_device_vector(handle, n_rows); + auto soar_labels = raft::make_device_vector(handle, n_rows); + + raft::update_device(dataset.data_handle(), h_dataset.data(), h_dataset.size(), stream); + raft::update_device(centroids.data_handle(), h_centroids.data(), h_centroids.size(), stream); + raft::update_device(labels.data_handle(), h_labels.data(), h_labels.size(), stream); + + cuvs::cluster::soar::params params; + cuvs::cluster::soar::predict(handle, + params, + raft::make_const_mdspan(dataset.view()), + raft::make_const_mdspan(centroids.view()), + raft::make_const_mdspan(labels.view()), + soar_labels.view()); + + std::vector result(n_rows); + raft::update_host(result.data(), soar_labels.data_handle(), n_rows, stream); + raft::resource::sync_stream(handle); + + ASSERT_EQ(expected, result); +} + +TEST(SoarTestErrors, RejectsMismatchedShapes) +{ + raft::resources handle; + + constexpr int64_t n_rows = 32, dim = 4, n_clusters = 8; + + auto dataset = raft::make_device_matrix(handle, n_rows, dim); + auto centroids = raft::make_device_matrix(handle, n_clusters, dim); + auto labels = raft::make_device_vector(handle, n_rows); + auto soar_labels = raft::make_device_vector(handle, n_rows); + + cuvs::cluster::soar::params params; + auto dataset_view = raft::make_const_mdspan(dataset.view()); + auto centroids_view = raft::make_const_mdspan(centroids.view()); + auto labels_view = raft::make_const_mdspan(labels.view()); + + // centroid dimension differs from the dataset dimension + auto narrow_centroids = raft::make_device_matrix(handle, n_clusters, dim - 1); + ASSERT_THROW(cuvs::cluster::soar::predict(handle, + params, + dataset_view, + raft::make_const_mdspan(narrow_centroids.view()), + labels_view, + soar_labels.view()), + raft::logic_error); + + // no centroids to choose from + auto no_centroids = raft::make_device_matrix(handle, 0, dim); + ASSERT_THROW(cuvs::cluster::soar::predict(handle, + params, + dataset_view, + raft::make_const_mdspan(no_centroids.view()), + labels_view, + soar_labels.view()), + raft::logic_error); + + // one primary label per row is required + auto short_labels = raft::make_device_vector(handle, n_rows - 1); + ASSERT_THROW(cuvs::cluster::soar::predict(handle, + params, + dataset_view, + centroids_view, + raft::make_const_mdspan(short_labels.view()), + soar_labels.view()), + raft::logic_error); + + // one output slot per row is required + auto short_output = raft::make_device_vector(handle, n_rows - 1); + ASSERT_THROW(cuvs::cluster::soar::predict( + handle, params, dataset_view, centroids_view, labels_view, short_output.view()), + raft::logic_error); +} + +} // namespace cuvs::cluster::soar diff --git a/examples/cpp/CMakeLists.txt b/examples/cpp/CMakeLists.txt index 91896dba56..4f39e4bb27 100644 --- a/examples/cpp/CMakeLists.txt +++ b/examples/cpp/CMakeLists.txt @@ -53,6 +53,7 @@ add_executable(IVF_FLAT_EXAMPLE src/ivf_flat_example.cu) add_executable(IVF_PQ_EXAMPLE src/ivf_pq_example.cu) add_executable(VAMANA_EXAMPLE src/vamana_example.cu) add_executable(SCANN_EXAMPLE src/scann_example.cu) +add_executable(SOAR_EXAMPLE src/soar_example.cu) # `$` is a generator expression that ensures that targets are # installed in a conda environment, if one exists @@ -87,3 +88,4 @@ target_link_libraries(IVF_PQ_EXAMPLE PRIVATE cuvs::cuvs $) target_link_libraries(VAMANA_EXAMPLE PRIVATE cuvs::cuvs $) target_link_libraries(SCANN_EXAMPLE PRIVATE cuvs::cuvs $) +target_link_libraries(SOAR_EXAMPLE PRIVATE cuvs::cuvs $) diff --git a/examples/cpp/src/soar_example.cu b/examples/cpp/src/soar_example.cu new file mode 100644 index 0000000000..8821a164de --- /dev/null +++ b/examples/cpp/src/soar_example.cu @@ -0,0 +1,102 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include + +#include +#include +#include +#include +#include +#include + +#include + +#include +#include + +/** Number of rows whose secondary cluster differs from their primary one. */ +int64_t count_spilled(raft::device_resources const& dev_resources, + raft::device_vector_view labels, + raft::device_vector_view soar_labels) +{ + auto h_labels = raft::make_host_vector(labels.extent(0)); + auto h_soar_labels = raft::make_host_vector(soar_labels.extent(0)); + auto stream = raft::resource::get_cuda_stream(dev_resources); + + raft::copy(h_labels.data_handle(), labels.data_handle(), labels.size(), stream); + raft::copy(h_soar_labels.data_handle(), soar_labels.data_handle(), soar_labels.size(), stream); + raft::resource::sync_stream(dev_resources, stream); + + int64_t n_spilled = 0; + for (int64_t i = 0; i < labels.extent(0); ++i) { + if (h_soar_labels(i) != h_labels(i)) { ++n_spilled; } + } + return n_spilled; +} + +void soar_predict_example(raft::device_resources const& dev_resources, + raft::device_matrix_view dataset, + raft::device_matrix_view centroids, + raft::device_vector_view labels) +{ + // Default lambda = 1. Larger values penalize secondary centroids whose residual is aligned with + // the primary residual, favoring more complementary assignments. + cuvs::cluster::soar::params params; + + auto soar_labels = raft::make_device_vector(dev_resources, dataset.extent(0)); + + cuvs::cluster::soar::predict( + dev_resources, params, dataset, centroids, labels, soar_labels.view()); + + // A row keeps its primary label when no other centroid is worth spilling to. + auto n_spilled = + count_spilled(dev_resources, labels, raft::make_const_mdspan(soar_labels.view())); + + std::cout << "Spilled " << n_spilled << " of " << dataset.extent(0) + << " rows to a secondary cluster" << std::endl; +} + +int main() +{ + raft::device_resources dev_resources; + + // Set pool memory resource with 1 GiB initial pool size. All allocations use the same pool. + rmm::mr::pool_memory_resource pool_mr(rmm::mr::get_current_device_resource_ref(), + 1024 * 1024 * 1024ull); + rmm::mr::set_current_device_resource(pool_mr); + + int64_t n_samples = 10000; + int64_t n_dim = 64; + int64_t n_clusters = 100; + + // Far fewer blobs than k-means clusters: the 100 learned partitions subdivide the 10 dense + // regions, creating internal partition boundaries. blob_labels is required by make_blobs but + // unused. + int64_t n_blobs = 10; + auto dataset = raft::make_device_matrix(dev_resources, n_samples, n_dim); + auto blob_labels = raft::make_device_vector(dev_resources, n_samples); + raft::random::make_blobs(dev_resources, dataset.view(), blob_labels.view(), n_blobs); + + auto dataset_view = raft::make_const_mdspan(dataset.view()); + + // SOAR needs centroids and a primary label per row, so k-means runs first. + cuvs::cluster::kmeans::balanced_params kmeans_params; + auto centroids = raft::make_device_matrix(dev_resources, n_clusters, n_dim); + auto labels = raft::make_device_vector(dev_resources, n_samples); + + cuvs::cluster::kmeans::fit(dev_resources, kmeans_params, dataset_view, centroids.view()); + cuvs::cluster::kmeans::predict(dev_resources, + kmeans_params, + dataset_view, + raft::make_const_mdspan(centroids.view()), + labels.view()); + + soar_predict_example(dev_resources, + dataset_view, + raft::make_const_mdspan(centroids.view()), + raft::make_const_mdspan(labels.view())); +} diff --git a/fern/docs.yml b/fern/docs.yml index 77ee0dae2a..4b9e0a17e0 100644 --- a/fern/docs.yml +++ b/fern/docs.yml @@ -274,6 +274,8 @@ navigation: path: "./pages/c_api/c-api-cluster-kmeans.md" - page: "Cluster Multi GPU Kmeans" path: "./pages/c_api/c-api-cluster-mg-kmeans.md" + - page: "Cluster Soar" + path: "./pages/c_api/c-api-cluster-soar.md" - page: "Core C API" path: "./pages/c_api/c-api-core-c-api.md" - page: "Core Dataset" @@ -333,6 +335,8 @@ navigation: path: "./pages/cpp_api/cpp-api-cluster-gmm.md" - page: "Cluster Kmeans" path: "./pages/cpp_api/cpp-api-cluster-kmeans.md" + - page: "Cluster Soar" + path: "./pages/cpp_api/cpp-api-cluster-soar.md" - page: "Cluster Spectral" path: "./pages/cpp_api/cpp-api-cluster-spectral.md" - page: "Core Bloom Filter" @@ -431,6 +435,8 @@ navigation: path: "./pages/python_api/python-api-cluster-kmeans.md" - page: "Cluster Multi GPU Kmeans" path: "./pages/python_api/python-api-cluster-mg-kmeans.md" + - page: "Cluster Soar" + path: "./pages/python_api/python-api-cluster-soar.md" - page: "Common" path: "./pages/python_api/python-api-common.md" - page: "Distance" diff --git a/fern/pages/c_api/c-api-cluster-soar.md b/fern/pages/c_api/c-api-cluster-soar.md new file mode 100644 index 0000000000..d4d2aa8760 --- /dev/null +++ b/fern/pages/c_api/c-api-cluster-soar.md @@ -0,0 +1,107 @@ +--- +slug: api-reference/c-api-cluster-soar +--- + +# Soar + +_Source header: `cuvs/cluster/soar.h`_ + +## SOAR hyperparameters + + +### cuvsSoarParams + +Hyper-parameters for SOAR assignment. + +```c +struct cuvsSoarParams { + float lambda; +}; +``` + +**Fields** + +| Name | Type | Description | +| --- | --- | --- | +| `lambda` | `float` | Weight of the projection of the secondary residual onto the primary residual in the SOAR loss. Larger values penalize secondary centroids whose residual is aligned with the primary residual, favoring complementary assignments. `0` reduces the loss to plain squared distance, which the primary centroid itself minimizes, so nothing is spilled. Default: 1.0. | + + +### cuvsSoarParamsCreate + +Allocate SOAR params, and populate with default values + +```c +cuvsError_t cuvsSoarParamsCreate(cuvsSoarParams_t* params); +``` + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `params` | out | [`cuvsSoarParams_t*`](/api-reference/c-api-cluster-soar#cuvssoarparams) | cuvsSoarParams_t to allocate | + +**Returns** + +[`cuvsError_t`](/api-reference/c-api-core-c-api#cuvserror-t) + + +### cuvsSoarParamsDestroy + +De-allocate SOAR params + +```c +cuvsError_t cuvsSoarParamsDestroy(cuvsSoarParams_t params); +``` + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `params` | in | [`cuvsSoarParams_t`](/api-reference/c-api-cluster-soar#cuvssoarparams) | cuvsSoarParams_t to de-allocate | + +**Returns** + +[`cuvsError_t`](/api-reference/c-api-core-c-api#cuvserror-t) + +## SOAR assignment + + +### cuvsSoarPredict + +Assign a secondary ("spilled") cluster to each row of the dataset. + +```c +cuvsError_t cuvsSoarPredict(cuvsResources_t res, +cuvsSoarParams_t params, +DLManagedTensor* dataset, +DLManagedTensor* centroids, +DLManagedTensor* labels, +DLManagedTensor* soar_labels); +``` + +SOAR (Spilling with Orthogonality-Amplified Residuals) picks, for each vector, a second centroid that complements the primary assignment instead of merely being the next-closest one. It minimizes the loss of Theorem 3.1 of https://arxiv.org/abs/2404.00774: for a vector `x` with primary residual `r = x - centroids[labels[i]]`, + +`score(c) = \|\|x - c\|\|^2 + lambda * (dot(r / \|\|r\|\|, x - c))^2` + +and `soar_labels[i]` is the centroid minimizing that score. Indexing a vector under both its primary and its secondary centroid improves recall for queries near a partition boundary. + +All tensors must be on device memory. `dataset` and `centroids` must be row-major float32. `labels` and `soar_labels` must have the same dtype, either uint32 or int32; int32 is accepted so that the output of `cuvsKMeansPredict` can be passed through without a conversion. + +The primary centroid is not excluded from the search, so `soar_labels[i] == labels[i]` is a possible (and meaningful) result: it says that no other centroid is worth spilling to, which is the common case for vectors in the interior of a cluster. + +Scratch memory scales as `n_rows * n_clusters * 4` bytes because scores against all centroids are materialized at once and are not tiled. Process the dataset in row batches to bound the peak device memory usage. + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `res` | in | [`cuvsResources_t`](/api-reference/c-api-core-c-api#cuvsresources-t) | opaque C handle | +| `params` | in | [`cuvsSoarParams_t`](/api-reference/c-api-cluster-soar#cuvssoarparams) | Parameters for SOAR assignment. | +| `dataset` | in | `DLManagedTensor*` | The dataset. The data must be in row-major format. [dim = n_rows x n_features] | +| `centroids` | in | `DLManagedTensor*` | Cluster centroids. The data must be in row-major format. [dim = n_clusters x n_features] | +| `labels` | in | `DLManagedTensor*` | Index of the primary cluster each row belongs to, as produced by k-means prediction. Every value must be in `[0, n_clusters)`. [len = n_rows] | +| `soar_labels` | out | `DLManagedTensor*` | Index of the secondary cluster each row is spilled to. [len = n_rows] | + +**Returns** + +[`cuvsError_t`](/api-reference/c-api-core-c-api#cuvserror-t) diff --git a/fern/pages/c_api/index.md b/fern/pages/c_api/index.md index cd70102764..ce2748219d 100644 --- a/fern/pages/c_api/index.md +++ b/fern/pages/c_api/index.md @@ -6,6 +6,7 @@ These pages are generated from the documented public headers in the cuVS source - [K-Means](/api-reference/c-api-cluster-kmeans) - [Multi-GPU K-Means](/api-reference/c-api-cluster-mg-kmeans) +- [Soar](/api-reference/c-api-cluster-soar) ## Common diff --git a/fern/pages/cpp_api/cpp-api-cluster-soar.md b/fern/pages/cpp_api/cpp-api-cluster-soar.md new file mode 100644 index 0000000000..466fcb789f --- /dev/null +++ b/fern/pages/cpp_api/cpp-api-cluster-soar.md @@ -0,0 +1,69 @@ +--- +slug: api-reference/cpp-api-cluster-soar +--- + +# Soar + +_Source header: `cuvs/cluster/soar.hpp`_ + +## SOAR hyperparameters + + +### cluster::soar::params + +Simple object to specify hyper-parameters for SOAR assignment. + +```cpp +struct params { + float lambda; +}; +``` + +**Fields** + +| Name | Type | Description | +| --- | --- | --- | +| `lambda` | `float` | Weight of the projection of the secondary residual onto the primary residual in the SOAR loss. Larger values penalize secondary centroids whose residual is aligned with the primary residual, favoring complementary assignments. `0` reduces the loss to plain squared distance, which the primary centroid itself minimizes, so nothing is spilled. Default: 1.0. | + +## SOAR assignment + + +### cluster::soar::predict + +Assign a secondary ("spilled") cluster to each row of the dataset. + +```cpp +void predict(raft::resources const& handle, +const soar::params& params, +raft::device_matrix_view dataset, +raft::device_matrix_view centroids, +raft::device_vector_view labels, +raft::device_vector_view soar_labels); +``` + +SOAR (Spilling with Orthogonality-Amplified Residuals) picks, for each vector, a second centroid that complements the primary assignment instead of merely being the next-closest one. It minimizes the loss of Theorem 3.1 of https://arxiv.org/abs/2404.00774: for a vector `x` with primary residual `r = x - centroids[labels[i]]`, + +`score(c) = \|\|x - c\|\|^2 + lambda * (dot(r / \|\|r\|\|, x - c))^2` + +and `soar_labels[i]` is the centroid minimizing that score. Indexing a vector under both its primary and its secondary centroid improves recall for queries near a partition boundary. + +Only float32 data and uint32 labels are supported. + +The primary centroid is not excluded from the search, so `soar_labels[i] == labels[i]` is a possible (and meaningful) result: it says that no other centroid is worth spilling to, which is the common case for vectors in the interior of a cluster. Callers that treat SOAR as a strictly second posting list should test for this case and skip those rows. + +Scratch memory scales as `n_rows * n_clusters * 4` bytes because scores against all centroids are materialized at once and are not tiled. Process the dataset in row batches to bound the peak device memory usage. + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `handle` | in | `raft::resources const&` | The raft handle. | +| `params` | in | [`const soar::params&`](/api-reference/cpp-api-cluster-soar#cluster-soar-params) | Parameters for SOAR assignment. | +| `dataset` | in | `raft::device_matrix_view` | The dataset. The data must be in row-major format. [dim = n_rows x n_features] | +| `centroids` | in | `raft::device_matrix_view` | Cluster centroids. The data must be in row-major format. [dim = n_clusters x n_features] | +| `labels` | in | `raft::device_vector_view` | Index of the primary cluster each row belongs to, as produced by k-means prediction. Every value must be in `[0, n_clusters)`. [len = n_rows] | +| `soar_labels` | out | `raft::device_vector_view` | Index of the secondary cluster each row is spilled to. [len = n_rows] | + +**Returns** + +`void` diff --git a/fern/pages/cpp_api/index.md b/fern/pages/cpp_api/index.md index 07115a5eca..9684276ed5 100644 --- a/fern/pages/cpp_api/index.md +++ b/fern/pages/cpp_api/index.md @@ -7,6 +7,7 @@ These pages are generated from the documented public headers in the cuVS source - [Agglomerative](/api-reference/cpp-api-cluster-agglomerative) - [Gmm](/api-reference/cpp-api-cluster-gmm) - [K-Means](/api-reference/cpp-api-cluster-kmeans) +- [Soar](/api-reference/cpp-api-cluster-soar) - [Spectral](/api-reference/cpp-api-cluster-spectral) ## Common diff --git a/fern/pages/python_api/index.md b/fern/pages/python_api/index.md index 897210d3ef..8245fb9ddb 100644 --- a/fern/pages/python_api/index.md +++ b/fern/pages/python_api/index.md @@ -6,6 +6,7 @@ These pages are generated from the Python and Cython sources under `python/cuvs/ - [Kmeans](/api-reference/python-api-cluster-kmeans) - [Kmeans](/api-reference/python-api-cluster-mg-kmeans) +- [Soar](/api-reference/python-api-cluster-soar) ## Common diff --git a/fern/pages/python_api/python-api-cluster-soar.md b/fern/pages/python_api/python-api-cluster-soar.md new file mode 100644 index 0000000000..5d93b5bc23 --- /dev/null +++ b/fern/pages/python_api/python-api-cluster-soar.md @@ -0,0 +1,115 @@ +--- +slug: api-reference/python-api-cluster-soar +--- + +# Soar + +_Python module: `cuvs.cluster.soar`_ + +## Params + +```python +cdef class Params +``` + +Hyper-parameters for SOAR assignment. + +**Parameters** + +| Name | Type | Description | +| --- | --- | --- | +| `lambda_` | `float` | Weight of the projection of the secondary residual onto the primary residual in the SOAR loss. Larger values penalize secondary centroids whose residual is aligned with the primary residual, favoring complementary assignments. ``0`` reduces the loss to plain squared distance, which the primary centroid itself minimizes, so nothing is spilled. Named with a trailing underscore because ``lambda`` is a Python keyword (default: 1.0). | + +**Constructor** + +```python +def __init__(self, *, lambda_=None) +``` + +**Members** + +| Name | Kind | +| --- | --- | +| `lambda_` | property | + +### lambda_ + +```python +def lambda_(self) +``` + +## predict + +`@auto_sync_resources` +`@auto_convert_output` + +```python +def predict(Params params, dataset, centroids, labels, soar_labels=None, resources=None) +``` + +Assign a secondary ("spilled") cluster to each row of the dataset. + +SOAR (Spilling with Orthogonality-Amplified Residuals) picks, for each +vector, a second centroid that complements the primary assignment instead +of merely being the next-closest one. It minimizes the loss of Theorem 3.1 +of https://arxiv.org/abs/2404.00774: for a vector ``x`` with primary +residual ``r = x - centroids[labels[i]]``, + +``score(c) = \|\|x - c\|\|^2 + lambda * (dot(r / \|\|r\|\|, x - c))^2`` + +and ``soar_labels[i]`` is the centroid minimizing that score. Indexing a +vector under both its primary and its secondary centroid improves recall +for queries near a partition boundary. + +The primary centroid is not excluded from the search, so +``soar_labels[i] == labels[i]`` is a possible (and meaningful) result: it +says that no other centroid is worth spilling to, which is the common case +for vectors in the interior of a cluster. Callers that treat SOAR as a +strictly second posting list should test for this case and skip those rows. + +Scratch memory scales as ``n_rows * n_clusters * 4`` bytes because scores +against all centroids are materialized at once and are not tiled. Process +the dataset in row batches to bound the peak device memory usage. + +**Parameters** + +| Name | Type | Description | +| --- | --- | --- | +| `params` | `Params` | Parameters for SOAR assignment. | +| `dataset` | `CUDA array interface compliant matrix, row major, float32` | shape (n_rows, n_features) | +| `centroids` | `CUDA array interface compliant matrix, row major, float32` | Cluster centroids, shape (n_clusters, n_features) | +| `labels` | `CUDA array interface compliant vector, uint32 or int32` | Index of the primary cluster each row belongs to, as produced by k-means prediction. Every value must be in ``[0, n_clusters)``. shape (n_rows,) | +| `soar_labels` | `Optional preallocated CUDA array interface vector to hold` | the output, shape (n_rows,). Must have the same dtype as ``labels``. When None, an array matching the dtype of ``labels`` is allocated. | +| `resources` | `cuvs.common.Resources, optional` | | + +**Returns** + +| Name | Type | Description | +| --- | --- | --- | +| `soar_labels` | `raft.device_ndarray` | Index of the secondary cluster each row is spilled to. | + +**Examples** + +```python +>>> import cupy as cp +>>> +>>> from cuvs.cluster.kmeans import KMeansParams +>>> from cuvs.cluster.kmeans import fit as kmeans_fit +>>> from cuvs.cluster.kmeans import predict as kmeans_predict +>>> from cuvs.cluster.soar import Params, predict +>>> +>>> n_samples = 5000 +>>> n_features = 50 +>>> n_clusters = 50 +>>> +>>> X = cp.random.random_sample((n_samples, n_features), +... dtype=cp.float32) +>>> +>>> # SOAR needs centroids and a primary label per row, so k-means runs +>>> # first +>>> kmeans_params = KMeansParams(n_clusters=n_clusters) +>>> centroids, inertia, n_iter = kmeans_fit(kmeans_params, X) +>>> labels, inertia = kmeans_predict(kmeans_params, X, centroids) +>>> +>>> soar_labels = predict(Params(), X, centroids, labels) +``` diff --git a/python/cuvs/cuvs/cluster/CMakeLists.txt b/python/cuvs/cuvs/cluster/CMakeLists.txt index 97f0a52f37..d07eba5af8 100644 --- a/python/cuvs/cuvs/cluster/CMakeLists.txt +++ b/python/cuvs/cuvs/cluster/CMakeLists.txt @@ -7,3 +7,4 @@ add_subdirectory(kmeans) add_subdirectory(mg) +add_subdirectory(soar) diff --git a/python/cuvs/cuvs/cluster/__init__.py b/python/cuvs/cuvs/cluster/__init__.py index 3de1f6fbc7..d91f4def4d 100644 --- a/python/cuvs/cuvs/cluster/__init__.py +++ b/python/cuvs/cuvs/cluster/__init__.py @@ -2,6 +2,6 @@ # SPDX-License-Identifier: Apache-2.0 -from cuvs.cluster import kmeans, mg +from cuvs.cluster import kmeans, mg, soar -__all__ = ["kmeans", "mg"] +__all__ = ["kmeans", "mg", "soar"] diff --git a/python/cuvs/cuvs/cluster/soar/CMakeLists.txt b/python/cuvs/cuvs/cluster/soar/CMakeLists.txt new file mode 100644 index 0000000000..ee68555622 --- /dev/null +++ b/python/cuvs/cuvs/cluster/soar/CMakeLists.txt @@ -0,0 +1,17 @@ +# ============================================================================= +# cmake-format: off +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# cmake-format: on +# ============================================================================= + +# Set the list of Cython files to build +set(cython_sources soar.pyx) +set(linked_libraries cuvs::cuvs cuvs::c_api) + +# Build all of the Cython targets +rapids_cython_create_modules( + CXX + SOURCE_FILES "${cython_sources}" + LINKED_LIBRARIES "${linked_libraries}" ASSOCIATED_TARGETS cuvs MODULE_PREFIX cluster_soar_ +) diff --git a/python/cuvs/cuvs/cluster/soar/__init__.py b/python/cuvs/cuvs/cluster/soar/__init__.py new file mode 100644 index 0000000000..5f478dd6ce --- /dev/null +++ b/python/cuvs/cuvs/cluster/soar/__init__.py @@ -0,0 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from .soar import Params, predict + +__all__ = ["Params", "predict"] diff --git a/python/cuvs/cuvs/cluster/soar/soar.pxd b/python/cuvs/cuvs/cluster/soar/soar.pxd new file mode 100644 index 0000000000..2351f69c78 --- /dev/null +++ b/python/cuvs/cuvs/cluster/soar/soar.pxd @@ -0,0 +1,26 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# cython: language_level=3 + +from cuvs.common.c_api cimport cuvsError_t, cuvsResources_t +from cuvs.common.cydlpack cimport DLManagedTensor + + +cdef extern from "cuvs/cluster/soar.h" nogil: + ctypedef struct cuvsSoarParams: + float lambda_ "lambda" + + ctypedef cuvsSoarParams* cuvsSoarParams_t + + cuvsError_t cuvsSoarParamsCreate(cuvsSoarParams_t* params) + + cuvsError_t cuvsSoarParamsDestroy(cuvsSoarParams_t params) + + cuvsError_t cuvsSoarPredict(cuvsResources_t res, + cuvsSoarParams_t params, + DLManagedTensor* dataset, + DLManagedTensor* centroids, + DLManagedTensor* labels, + DLManagedTensor* soar_labels) diff --git a/python/cuvs/cuvs/cluster/soar/soar.pyx b/python/cuvs/cuvs/cluster/soar/soar.pyx new file mode 100644 index 0000000000..36fc43004d --- /dev/null +++ b/python/cuvs/cuvs/cluster/soar/soar.pyx @@ -0,0 +1,171 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# cython: language_level=3 + +import numpy as np + +from cuvs.common cimport cydlpack + +from pylibraft.common import auto_convert_output, device_ndarray +from pylibraft.common.cai_wrapper import wrap_array +from pylibraft.common.interruptible import cuda_interruptible + +from cuvs.common.exceptions import check_cuvs +from cuvs.common.resources import auto_sync_resources +from cuvs.neighbors.common import _check_input_array + +LABEL_DTYPES = [np.dtype("uint32"), np.dtype("int32")] + + +cdef class Params: + """ + Hyper-parameters for SOAR assignment. + + Parameters + ---------- + lambda_ : float + Weight of the projection of the secondary residual onto the primary + residual in the SOAR loss. Larger values penalize secondary centroids + whose residual is aligned with the primary residual, favoring + complementary assignments. ``0`` reduces the loss to plain squared + distance, which the primary centroid itself minimizes, so nothing is + spilled. Named with a trailing underscore because ``lambda`` is a + Python keyword (default: 1.0). + """ + + cdef cuvsSoarParams* params + + def __cinit__(self): + check_cuvs(cuvsSoarParamsCreate(&self.params)) + + def __dealloc__(self): + check_cuvs(cuvsSoarParamsDestroy(self.params)) + + def __init__(self, *, lambda_=None): + if lambda_ is not None: + self.params.lambda_ = lambda_ + + @property + def lambda_(self): + return self.params.lambda_ + + +@auto_sync_resources +@auto_convert_output +def predict(Params params, dataset, centroids, labels, soar_labels=None, + resources=None): + """ + Assign a secondary ("spilled") cluster to each row of the dataset. + + SOAR (Spilling with Orthogonality-Amplified Residuals) picks, for each + vector, a second centroid that complements the primary assignment instead + of merely being the next-closest one. It minimizes the loss of Theorem 3.1 + of https://arxiv.org/abs/2404.00774: for a vector ``x`` with primary + residual ``r = x - centroids[labels[i]]``, + + ``score(c) = ||x - c||^2 + lambda * (dot(r / ||r||, x - c))^2`` + + and ``soar_labels[i]`` is the centroid minimizing that score. Indexing a + vector under both its primary and its secondary centroid improves recall + for queries near a partition boundary. + + The primary centroid is not excluded from the search, so + ``soar_labels[i] == labels[i]`` is a possible (and meaningful) result: it + says that no other centroid is worth spilling to, which is the common case + for vectors in the interior of a cluster. Callers that treat SOAR as a + strictly second posting list should test for this case and skip those rows. + + Scratch memory scales as ``n_rows * n_clusters * 4`` bytes because scores + against all centroids are materialized at once and are not tiled. Process + the dataset in row batches to bound the peak device memory usage. + + Parameters + ---------- + params : Params + Parameters for SOAR assignment. + dataset : CUDA array interface compliant matrix, row major, float32 + shape (n_rows, n_features) + centroids : CUDA array interface compliant matrix, row major, float32 + Cluster centroids, shape (n_clusters, n_features) + labels : CUDA array interface compliant vector, uint32 or int32 + Index of the primary cluster each row belongs to, as produced by + k-means prediction. Every value must be in ``[0, n_clusters)``. + shape (n_rows,) + soar_labels : Optional preallocated CUDA array interface vector to hold + the output, shape (n_rows,). Must have the same dtype as ``labels``. + When None, an array matching the dtype of ``labels`` is allocated. + {resources_docstring} + + Returns + ------- + soar_labels : raft.device_ndarray + Index of the secondary cluster each row is spilled to. + + Examples + -------- + + >>> import cupy as cp + >>> + >>> from cuvs.cluster.kmeans import KMeansParams + >>> from cuvs.cluster.kmeans import fit as kmeans_fit + >>> from cuvs.cluster.kmeans import predict as kmeans_predict + >>> from cuvs.cluster.soar import Params, predict + >>> + >>> n_samples = 5000 + >>> n_features = 50 + >>> n_clusters = 50 + >>> + >>> X = cp.random.random_sample((n_samples, n_features), + ... dtype=cp.float32) + >>> + >>> # SOAR needs centroids and a primary label per row, so k-means runs + >>> # first + >>> kmeans_params = KMeansParams(n_clusters=n_clusters) + >>> centroids, inertia, n_iter = kmeans_fit(kmeans_params, X) + >>> labels, inertia = kmeans_predict(kmeans_params, X, centroids) + >>> + >>> soar_labels = predict(Params(), X, centroids, labels) + """ + + dataset_ai = wrap_array(dataset) + _check_input_array(dataset_ai, [np.dtype("float32")]) + + centroids_ai = wrap_array(centroids) + _check_input_array(centroids_ai, [np.dtype("float32")], + exp_cols=dataset_ai.shape[1]) + + labels_ai = wrap_array(labels) + _check_input_array(labels_ai, LABEL_DTYPES, + exp_rows=dataset_ai.shape[0]) + + if soar_labels is None: + soar_labels = device_ndarray.empty((dataset_ai.shape[0],), + dtype=labels_ai.dtype) + + soar_labels_ai = wrap_array(soar_labels) + _check_input_array(soar_labels_ai, [labels_ai.dtype], + exp_rows=dataset_ai.shape[0]) + + cdef cydlpack.DLManagedTensor* dataset_dlpack = \ + cydlpack.dlpack_c(dataset_ai) + cdef cydlpack.DLManagedTensor* centroids_dlpack = \ + cydlpack.dlpack_c(centroids_ai) + cdef cydlpack.DLManagedTensor* labels_dlpack = \ + cydlpack.dlpack_c(labels_ai) + cdef cydlpack.DLManagedTensor* soar_labels_dlpack = \ + cydlpack.dlpack_c(soar_labels_ai) + + cdef cuvsResources_t res = resources.get_c_obj() + + with cuda_interruptible(): + check_cuvs(cuvsSoarPredict( + res, + params.params, + dataset_dlpack, + centroids_dlpack, + labels_dlpack, + soar_labels_dlpack)) + + return soar_labels diff --git a/python/cuvs/cuvs/tests/test_soar.py b/python/cuvs/cuvs/tests/test_soar.py new file mode 100644 index 0000000000..bdd3151a2b --- /dev/null +++ b/python/cuvs/cuvs/tests/test_soar.py @@ -0,0 +1,192 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + +import numpy as np +import pytest +from pylibraft.common import device_ndarray + +from cuvs.cluster.soar import Params, predict + +# Deliberately not the default of 1.0, so that the end-to-end test fails if +# lambda is dropped anywhere between Params and the C++ implementation. +LAMBDA = 2.0 + + +def nearest_centroids(dataset, centroids): + """Primary assignment, i.e. what k-means prediction would produce.""" + distances = ((dataset[:, None, :] - centroids[None, :, :]) ** 2).sum( + axis=-1 + ) + return np.argmin(distances, axis=1) + + +def soar_scores(dataset, centroids, labels, lambda_): + """ + ``||x - c||^2 + lambda * (dot(r / ||r||, x - c))^2`` for every row and + centroid, computed in float64 as the reference for the device kernel. + """ + dataset = dataset.astype(np.float64) + centroids = centroids.astype(np.float64) + + residual = dataset - centroids[labels] + unit_residual = residual / np.linalg.norm(residual, axis=1, keepdims=True) + + diff = dataset[:, None, :] - centroids[None, :, :] + squared_distance = (diff**2).sum(axis=-1) + projection = (diff * unit_residual[:, None, :]).sum(axis=-1) + + return squared_distance + lambda_ * projection**2 + + +def make_inputs(n_rows=256, n_features=8, n_clusters=16, seed=0): + rng = np.random.default_rng(seed) + dataset = rng.uniform(-1, 1, (n_rows, n_features)).astype(np.float32) + centroids = rng.uniform(-1, 1, (n_clusters, n_features)).astype(np.float32) + labels = nearest_centroids(dataset, centroids) + return dataset, centroids, labels + + +def test_params_defaults(): + assert Params().lambda_ == pytest.approx(1.0) + + +def test_params_custom(): + assert Params(lambda_=2.5).lambda_ == pytest.approx(2.5) + + +@pytest.mark.parametrize("n_rows", [256, 1000]) +@pytest.mark.parametrize("n_features", [8, 37]) +@pytest.mark.parametrize("n_clusters", [16, 100]) +def test_predict_matches_host_reference(n_rows, n_features, n_clusters): + dataset, centroids, labels = make_inputs(n_rows, n_features, n_clusters) + labels = labels.astype(np.uint32) + + soar_labels = predict( + Params(lambda_=LAMBDA), + device_ndarray(dataset), + device_ndarray(centroids), + device_ndarray(labels), + ).copy_to_host() + + assert soar_labels.dtype == labels.dtype + assert soar_labels.shape == labels.shape + assert np.all(soar_labels < centroids.shape[0]) + + # Guards against the assertion below going vacuous: if the fixture ever + # stopped spilling, simply echoing the primary labels would satisfy it. + assert np.count_nonzero(soar_labels != labels) > 0 + + # Compare losses rather than ids so the test is not fragile when two + # centroids tie. + scores = soar_scores(dataset, centroids, labels, LAMBDA) + achieved = np.take_along_axis( + scores, soar_labels.astype(np.int64)[:, None], axis=1 + ).squeeze(1) + + np.testing.assert_allclose(achieved, scores.min(axis=1), rtol=1e-4) + + +def test_predict_int32_matches_uint32(): + dataset, centroids, labels = make_inputs() + + def run(dtype): + return predict( + Params(lambda_=LAMBDA), + device_ndarray(dataset), + device_ndarray(centroids), + device_ndarray(labels.astype(dtype)), + ).copy_to_host() + + np.testing.assert_array_equal( + run(np.uint32).astype(np.int64), run(np.int32).astype(np.int64) + ) + + +def test_predict_with_preallocated_output(): + dataset, centroids, labels = make_inputs() + labels = labels.astype(np.uint32) + + out = device_ndarray(np.zeros(dataset.shape[0], dtype=np.uint32)) + returned = predict( + Params(lambda_=LAMBDA), + device_ndarray(dataset), + device_ndarray(centroids), + device_ndarray(labels), + soar_labels=out, + ) + + assert returned is out + expected = predict( + Params(lambda_=LAMBDA), + device_ndarray(dataset), + device_ndarray(centroids), + device_ndarray(labels), + ).copy_to_host() + np.testing.assert_array_equal(out.copy_to_host(), expected) + + +def test_predict_input_validation(): + dataset, centroids, labels = make_inputs() + labels = labels.astype(np.uint32) + n_rows = dataset.shape[0] + + params = Params() + dataset_d = device_ndarray(dataset) + centroids_d = device_ndarray(centroids) + labels_d = device_ndarray(labels) + + with pytest.raises(TypeError, match="dtype float64"): + predict( + params, + device_ndarray(dataset.astype(np.float64)), + device_ndarray(centroids.astype(np.float64)), + labels_d, + ) + + with pytest.raises(TypeError, match="dtype int64"): + predict( + params, + dataset_d, + centroids_d, + device_ndarray(labels.astype(np.int64)), + ) + + with pytest.raises(ValueError, match="Incorrect number of columns"): + predict( + params, + dataset_d, + device_ndarray(np.ascontiguousarray(centroids[:, :-1])), + labels_d, + ) + + with pytest.raises(ValueError, match="Incorrect number of rows"): + predict(params, dataset_d, centroids_d, device_ndarray(labels[:-1])) + + with pytest.raises(ValueError, match="Incorrect number of rows"): + predict( + params, + dataset_d, + centroids_d, + labels_d, + soar_labels=device_ndarray(np.zeros(n_rows - 1, dtype=np.uint32)), + ) + + # int32 is a valid label dtype, but the output must match the dtype of + # `labels`, so a uint32/int32 pair is still rejected. + with pytest.raises(TypeError, match="dtype int32"): + predict( + params, + dataset_d, + centroids_d, + labels_d, + soar_labels=device_ndarray(np.zeros(n_rows, dtype=np.int32)), + ) + + with pytest.raises(ValueError, match="Row major"): + predict( + params, + device_ndarray(np.asfortranarray(dataset)), + centroids_d, + labels_d, + )