Add integration for train only scenario

ibhati · ibhati · commit dbef61e2ece2 · 2025-11-13T17:13:35.000-08:00
diff --git a/include/svs/index/ivf/common.h b/include/svs/index/ivf/common.h
@@ -404,7 +404,9 @@ void centroid_assignment(
                     data.dimensions()
                 );
             }
-            if constexpr (std::is_same_v<Distance, distance::DistanceIP>) {
+            if constexpr (std::is_same_v<
+                              std::remove_cvref_t<Distance>,
+                              distance::DistanceIP>) {
                 for (auto i : indices) {
                     auto nearest =
                         type_traits::sentinel_v<Neighbor<size_t>, std::greater<>>;
@@ -414,7 +416,9 @@ void centroid_assignment(
                     }
                     assignments[batch_range.start() + i] = nearest.id();
                 }
-            } else if constexpr (std::is_same_v<Distance, distance::DistanceL2>) {
+            } else if constexpr (std::is_same_v<
+                                     std::remove_cvref_t<Distance>,
+                                     distance::DistanceL2>) {
                 for (auto i : indices) {
                     auto nearest = type_traits::sentinel_v<Neighbor<size_t>, std::less<>>;
                     auto dists = matmul_results.get_datum(i);
@@ -563,13 +567,13 @@ auto kmeans_training(
     auto training_timer = timer.push_back("Kmeans training");
     data::SimpleData<float> centroids_fp32 = convert_data<float>(centroids, threadpool);
 
-    if constexpr (std::is_same_v<Distance, distance::DistanceIP>) {
+    if constexpr (std::is_same_v<std::remove_cvref_t<Distance>, distance::DistanceIP>) {
         normalize_centroids(centroids_fp32, threadpool, timer);
     }
 
     auto assignments = std::vector<size_t>(data.size());
     std::vector<float> data_norm;
-    if constexpr (std::is_same_v<Distance, distance::DistanceL2>) {
+    if constexpr (std::is_same_v<std::remove_cvref_t<Distance>, distance::DistanceL2>) {
         generate_norms(data, data_norm, threadpool);
     }
     std::vector<float> centroids_norm;
@@ -578,7 +582,7 @@ auto kmeans_training(
         auto iter_timer = timer.push_back("iteration");
         auto batchsize = parameters.minibatch_size_;
         auto num_batches = lib::div_round_up(data.size(), batchsize);
-        if constexpr (std::is_same_v<Distance, distance::DistanceL2>) {
+        if constexpr (std::is_same_v<std::remove_cvref_t<Distance>, distance::DistanceL2>) {
             generate_norms(centroids_fp32, centroids_norm, threadpool);
         }
 
@@ -611,7 +615,7 @@ auto kmeans_training(
 
         centroid_split(data, centroids_fp32, counts, rng, threadpool, timer);
 
-        if constexpr (std::is_same_v<Distance, distance::DistanceIP>) {
+        if constexpr (std::is_same_v<std::remove_cvref_t<Distance>, distance::DistanceIP>) {
             normalize_centroids(centroids_fp32, threadpool, timer);
         }
     }
@@ -723,7 +727,7 @@ data::SimpleData<BuildType> init_centroids(
 template <typename Distance, typename Data, threads::ThreadPool Pool>
 std::vector<float> maybe_compute_norms(const Data& data, Pool& threadpool) {
     std::vector<float> norms;
-    if constexpr (std::is_same_v<Distance, distance::DistanceL2>) {
+    if constexpr (std::is_same_v<std::remove_cvref_t<Distance>, distance::DistanceL2>) {
         generate_norms(data, norms, threadpool);
     }
     return norms;
@@ -849,15 +853,15 @@ void search_centroids(
 ) {
     unsigned int count = 0;
     buffer.clear();
-    if constexpr (std::is_same_v<Dist, distance::DistanceIP>) {
+    if constexpr (std::is_same_v<std::remove_cvref_t<Dist>, distance::DistanceIP>) {
         for (size_t j = 0; j < num_threads; j++) {
             auto distance = matmul_results[j].get_datum(query_id);
             for (size_t k = 0; k < distance.size(); k++) {
                 buffer.insert({count, distance[k]});
                 count++;
             }
         }
-    } else if constexpr (std::is_same_v<Dist, distance::DistanceL2>) {
+    } else if constexpr (std::is_same_v<std::remove_cvref_t<Dist>, distance::DistanceL2>) {
         float query_norm = distance::norm_square(query);
         for (size_t j = 0; j < num_threads; j++) {
             auto distance = matmul_results[j].get_datum(query_id);
diff --git a/include/svs/index/ivf/index.h b/include/svs/index/ivf/index.h
@@ -320,7 +320,7 @@ class IVFIndex {
 
     void initialize_distance_metadata() {
         // Precalculate centroid norms for L2 distance
-        if constexpr (std::is_same_v<Dist, distance::DistanceL2>) {
+        if constexpr (std::is_same_v<std::remove_cvref_t<Dist>, distance::DistanceL2>) {
             centroids_norm_.reserve(centroids_.size());
             for (size_t i = 0; i < centroids_.size(); i++) {
                 centroids_norm_.push_back(distance::norm_square(centroids_.get_datum(i)));
diff --git a/tests/integration/ivf/index_build.cpp b/tests/integration/ivf/index_build.cpp
@@ -17,6 +17,10 @@
 // svs
 #include "svs/core/data/simple.h"
 #include "svs/core/recall.h"
+#include "svs/index/ivf/clustering.h"
+#include "svs/index/ivf/common.h"
+#include "svs/index/ivf/hierarchical_kmeans.h"
+#include "svs/lib/float16.h"
 #include "svs/lib/timing.h"
 #include "svs/orchestrators/ivf.h"
 
@@ -99,6 +103,96 @@ void test_build(const Distance& distance, size_t num_inner_threads = 1) {
     }
 }
 
+template <typename T, typename Distance>
+void test_build_train_only(const Distance& distance, size_t num_inner_threads = 1) {
+    const double epsilon = 0.06; // Wider tolerance for train_only workflow
+    const auto queries = svs::data::SimpleData<float>::load(test_dataset::query_file());
+    CATCH_REQUIRE(svs_test::prepare_temp_directory());
+    size_t num_threads = 2;
+
+    auto expected_result = test_dataset::ivf::expected_build_results(
+        svs::distance_type_v<Distance>, svsbenchmark::Uncompressed(svs::datatype_v<T>)
+    );
+
+    // Load data
+    auto data = svs::data::SimpleData<float>::load(test_dataset::data_svs_file());
+    auto threadpool = svs::threads::as_threadpool(num_threads);
+    auto parameters = expected_result.build_parameters_.value();
+
+    // Step 1: Use train_only mode to get centroids
+    svs::data::SimpleData<T> centroids_train;
+    std::vector<std::vector<uint32_t>> clusters_train;
+    fmt::print(
+        "Starting Train-Only Mode Clustering with {} centroids\n", parameters.num_centroids_
+    );
+
+    if (parameters.is_hierarchical_) {
+        fmt::print("Using Hierarchical KMeans Clustering\n");
+        std::tie(centroids_train, clusters_train) =
+            svs::index::ivf::hierarchical_kmeans_clustering<T>(
+                parameters,
+                data,
+                distance,
+                threadpool,
+                svs::lib::Type<uint32_t>(),
+                svs::logging::get(),
+                true // train_only = true
+            );
+    } else {
+        std::tie(centroids_train, clusters_train) = svs::index::ivf::kmeans_clustering<T>(
+            parameters,
+            data,
+            distance,
+            threadpool,
+            svs::lib::Type<uint32_t>(),
+            svs::logging::get(),
+            true // train_only = true
+        );
+    }
+
+    fmt::print("Train-Only Mode - Obtained {} centroids\n", centroids_train.size());
+
+    // Step 2: Assign data to clusters using cluster_assignment
+    auto clusters = svs::index::ivf::cluster_assignment<T>(
+        data,
+        centroids_train,
+        distance,
+        threadpool,
+        10'000, // minibatch_size
+        svs::lib::Type<uint32_t>()
+    );
+
+    // Step 3: Create clustering and assemble index
+    svs::index::ivf::Clustering clustering(std::move(centroids_train), std::move(clusters));
+
+    auto index = svs::IVF::assemble_from_clustering<float>(
+        std::move(clustering), std::move(data), distance, num_threads, num_inner_threads
+    );
+
+    // Test the index with the same expected results
+    auto groundtruth = test_dataset::load_groundtruth(svs::distance_type_v<Distance>);
+    for (const auto& expected : expected_result.config_and_recall_) {
+        auto these_queries = test_dataset::get_test_set(queries, expected.num_queries_);
+        auto these_groundtruth =
+            test_dataset::get_test_set(groundtruth, expected.num_queries_);
+        index.set_search_parameters(expected.search_parameters_);
+        auto results = index.search(these_queries, expected.num_neighbors_);
+        double recall = svs::k_recall_at_n(
+            these_groundtruth, results, expected.num_neighbors_, expected.recall_k_
+        );
+
+        fmt::print(
+            "Train-Only Mode - n_probes: {}, Expected Recall: {}, Actual Recall: {}\n",
+            index.get_search_parameters().n_probes_,
+            expected.recall_,
+            recall
+        );
+        // Just check that recall is reasonable (within wider tolerance)
+        CATCH_REQUIRE(recall > expected.recall_ - epsilon);
+        CATCH_REQUIRE(recall < expected.recall_ + epsilon);
+    }
+}
+
 } // namespace
 
 CATCH_TEST_CASE("IVF Build/Clustering", "[integration][build][ivf]") {
@@ -113,3 +207,11 @@ CATCH_TEST_CASE("IVF Build/Clustering", "[integration][build][ivf]") {
     // test_build<svs::BFloat16>(svs::DistanceL2(), 4);
     // test_build<svs::BFloat16>(svs::DistanceIP(), 4);
 }
+
+CATCH_TEST_CASE("IVF Build/Clustering", "[integration][build][ivf][train_only]") {
+    test_build_train_only<float>(svs::DistanceL2());
+    test_build_train_only<svs::Float16>(svs::DistanceIP());
+
+    test_build_train_only<svs::BFloat16>(svs::DistanceL2());
+    test_build_train_only<svs::BFloat16>(svs::DistanceIP());
+}
