faiss 0.5.3 → 0.6.1

data/vendor/faiss/faiss/impl/AdditiveQuantizer.cpp

@@ -21,6 +21,7 @@
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/LocalSearchQuantizer.h>
 #include <faiss/impl/ResidualQuantizer.h>
+#include <faiss/impl/simd_dispatch.h>
 #include <faiss/utils/Heap.h>
 #include <faiss/utils/distances.h>
 #include <faiss/utils/hamming.h>
@@ -47,13 +48,13 @@ int sgemm_(
 namespace faiss {
 
 AdditiveQuantizer::AdditiveQuantizer(
-        size_t d,
-        const std::vector<size_t>& nbits,
-        Search_type_t search_type)
-        : Quantizer(d),
-          M(nbits.size()),
-          nbits(nbits),
-          search_type(search_type) {
+        size_t d_in,
+        const std::vector<size_t>& nbits_in,
+        Search_type_t search_type_in)
+        : Quantizer(d_in),
+          M(nbits_in.size()),
+          nbits(nbits_in),
+          search_type(search_type_in) {
     set_derived_values();
 }
 
@@ -64,10 +65,12 @@ void AdditiveQuantizer::set_derived_values() {
     tot_bits = 0;
     only_8bit = true;
     codebook_offsets.resize(M + 1, 0);
-    for (int i = 0; i < M; i++) {
+    for (size_t i = 0; i < M; i++) {
         int nbit = nbits[i];
-        size_t k = 1 << nbit;
-        codebook_offsets[i + 1] = codebook_offsets[i] + k;
+        FAISS_CHECK_RANGE(nbit, 0, 31);
+        size_t k = (size_t)1 << nbit;
+        codebook_offsets[i + 1] =
+                add_no_overflow(codebook_offsets[i], k, "codebook_offsets");
         tot_bits += nbit;
         if (nbit != 0) {
             only_8bit = false;
@@ -104,7 +107,7 @@ void AdditiveQuantizer::set_derived_values() {
 void AdditiveQuantizer::train_norm(size_t n, const float* norms) {
     norm_min = HUGE_VALF;
     norm_max = -HUGE_VALF;
-    for (idx_t i = 0; i < n; i++) {
+    for (size_t i = 0; i < n; i++) {
         if (norms[i] < norm_min) {
             norm_min = norms[i];
         }
@@ -153,21 +156,42 @@ void AdditiveQuantizer::train_norm(size_t n, const float* norms) {
 
 void AdditiveQuantizer::compute_codebook_tables() {
     centroid_norms.resize(total_codebook_size);
+    FAISS_THROW_IF_NOT_FMT(
+            codebooks.size() >=
+                    mul_no_overflow(
+                            total_codebook_size, d, "codebooks validation"),
+            "codebooks size %zd too small for total_codebook_size=%zd * d=%zd",
+            codebooks.size(),
+            total_codebook_size,
+            d);
     fvec_norms_L2sqr(
             centroid_norms.data(), codebooks.data(), d, total_codebook_size);
     size_t cross_table_size = 0;
-    for (int m = 0; m < M; m++) {
+    for (size_t m = 0; m < M; m++) {
+        FAISS_CHECK_RANGE(nbits[m], 0, 31);
         size_t K = (size_t)1 << nbits[m];
-        cross_table_size += K * codebook_offsets[m];
+        size_t product =
+                mul_no_overflow(K, codebook_offsets[m], "cross_table_size");
+        cross_table_size = add_no_overflow(
+                cross_table_size, product, "cross_table_size accumulation");
     }
     codebook_cross_products.resize(cross_table_size);
     size_t ofs = 0;
-    for (int m = 1; m < M; m++) {
+    for (size_t m = 1; m < M; m++) {
         FINTEGER ki = (size_t)1 << nbits[m];
         FINTEGER kk = codebook_offsets[m];
         FINTEGER di = d;
         float zero = 0, one = 1;
-        assert(ofs + ki * kk <= cross_table_size);
+        size_t step_size = (size_t)ki * (size_t)kk;
+        FAISS_THROW_IF_NOT_FMT(
+                add_no_overflow(ofs, step_size, "cross product table offset") <=
+                        cross_table_size,
+                "cross product table overflow at step %zd: "
+                "%zd + %zd > %zd",
+                m,
+                ofs,
+                step_size,
+                cross_table_size);
         sgemm_("Transposed",
                "Not transposed",
                &ki,
@@ -181,7 +205,7 @@ void AdditiveQuantizer::compute_codebook_tables() {
                &zero,
                codebook_cross_products.data() + ofs,
                &ki);
-        ofs += ki * kk;
+        ofs += step_size;
     }
 }
 
@@ -276,11 +300,12 @@ void AdditiveQuantizer::pack_codes(
             norms = norm_buf.data();
         }
     }
+    int64_t n_signed = n;
 #pragma omp parallel for if (n > 1000)
-    for (int64_t i = 0; i < n; i++) {
+    for (int64_t i = 0; i < n_signed; i++) {
         const int32_t* codes1 = codes + i * ld_codes;
         BitstringWriter bsw(packed_codes + i * code_size, code_size);
-        for (int m = 0; m < M; m++) {
+        for (size_t m = 0; m < M; m++) {
             bsw.write(codes1[m], nbits[m]);
         }
         if (norm_bits != 0) {
@@ -293,12 +318,13 @@ void AdditiveQuantizer::decode(const uint8_t* code, float* x, size_t n) const {
     FAISS_THROW_IF_NOT_MSG(
             is_trained, "The additive quantizer is not trained yet.");
 
+    int64_t n_signed = n;
     // standard additive quantizer decoding
 #pragma omp parallel for if (n > 100)
-    for (int64_t i = 0; i < n; i++) {
+    for (int64_t i = 0; i < n_signed; i++) {
         BitstringReader bsr(code + i * code_size, code_size);
         float* xi = x + i * d;
-        for (int m = 0; m < M; m++) {
+        for (size_t m = 0; m < M; m++) {
             int idx = bsr.read(nbits[m]);
             const float* c = codebooks.data() + d * (codebook_offsets[m] + idx);
             if (m == 0) {
@@ -322,12 +348,13 @@ void AdditiveQuantizer::decode_unpacked(
         ld_codes = M;
     }
 
+    int64_t n_signed = n;
     // standard additive quantizer decoding
 #pragma omp parallel for if (n > 1000)
-    for (int64_t i = 0; i < n; i++) {
+    for (int64_t i = 0; i < n_signed; i++) {
         const int32_t* codesi = code + i * ld_codes;
         float* xi = x + i * d;
-        for (int m = 0; m < M; m++) {
+        for (size_t m = 0; m < M; m++) {
             int idx = codesi[m];
             const float* c = codebooks.data() + d * (codebook_offsets[m] + idx);
             if (m == 0) {
@@ -347,20 +374,23 @@ AdditiveQuantizer::~AdditiveQuantizer() {}
 
 void AdditiveQuantizer::compute_centroid_norms(float* norms) const {
     size_t ntotal = (size_t)1 << tot_bits;
+    int64_t ntotal_signed = ntotal;
     // TODO: make tree of partial sums
+    with_simd_level([&]<SIMDLevel SL>() {
 #pragma omp parallel
-    {
-        std::vector<float> tmp(d);
+        {
+            std::vector<float> tmp(d);
 #pragma omp for
-        for (int64_t i = 0; i < ntotal; i++) {
-            decode_64bit(i, tmp.data());
-            norms[i] = fvec_norm_L2sqr(tmp.data(), d);
+            for (int64_t i = 0; i < ntotal_signed; i++) {
+                decode_64bit(i, tmp.data());
+                norms[i] = fvec_norm_L2sqr<SL>(tmp.data(), d);
+            }
         }
-    }
+    });
 }
 
 void AdditiveQuantizer::decode_64bit(idx_t bits, float* xi) const {
-    for (int m = 0; m < M; m++) {
+    for (size_t m = 0; m < M; m++) {
         idx_t idx = bits & (((size_t)1 << nbits[m]) - 1);
         bits >>= nbits[m];
         const float* c = codebooks.data() + d * (codebook_offsets[m] + idx);
@@ -410,7 +440,7 @@ void compute_inner_prod_with_LUT(
         const float* LUT,
         float* ips) {
     size_t prev_size = 1;
-    for (int m = 0; m < aq.M; m++) {
+    for (size_t m = 0; m < aq.M; m++) {
         const float* LUTm = LUT + aq.codebook_offsets[m];
         int nb = aq.nbits[m];
         size_t nc = (size_t)1 << nb;
@@ -483,7 +513,7 @@ void AdditiveQuantizer::knn_centroids_L2(
             // ||x - y||^2 = ||x||^2 + ||y||^2 - 2 * <x,y>
 
             maxheap_heapify(k, distances_i, labels_i);
-            for (idx_t j = 0; j < ntotal; j++) {
+            for (size_t j = 0; j < ntotal; j++) {
                 float disj = q_norms[i] + norms[j] - 2 * dis[j];
                 if (disj < distances_i[0]) {
                     heap_replace_top<CMax<float, int64_t>>(
@@ -506,7 +536,7 @@ float accumulate_IPs(
         BitstringReader& bs,
         const float* LUT) {
     float accu = 0;
-    for (int m = 0; m < aq.M; m++) {
+    for (size_t m = 0; m < aq.M; m++) {
         size_t nbit = aq.nbits[m];
         int idx = bs.read(nbit);
         accu += LUT[idx];
@@ -519,7 +549,7 @@ float compute_norm_from_LUT(const AdditiveQuantizer& aq, BitstringReader& bs) {
     float accu = 0;
     std::vector<int> idx(aq.M);
     const float* c = aq.codebook_cross_products.data();
-    for (int m = 0; m < aq.M; m++) {
+    for (size_t m = 0; m < aq.M; m++) {
         size_t nbit = aq.nbits[m];
         int i = bs.read(nbit);
         size_t K = 1 << nbit;
@@ -527,7 +557,7 @@ float compute_norm_from_LUT(const AdditiveQuantizer& aq, BitstringReader& bs) {
 
         accu += aq.centroid_norms[aq.codebook_offsets[m] + i];
 
-        for (int l = 0; l < m; l++) {
+        for (size_t l = 0; l < m; l++) {
             int j = idx[l];
             accu += 2 * c[j * K + i];
             c += (1 << aq.nbits[l]) * K;

data/vendor/faiss/faiss/impl/AdditiveQuantizer.h

@@ -83,6 +83,7 @@ struct AdditiveQuantizer : Quantizer {
         ST_norm_lsq2x4, ///< use a 2x4 bits lsq as norm quantizer (for fast
                         ///< scan)
         ST_norm_rq2x4,  ///< use a 2x4 bits rq as norm quantizer (for fast scan)
+        ST_count
     };
 
     AdditiveQuantizer(

data/vendor/faiss/faiss/impl/AuxIndexStructures.cpp

@@ -20,10 +20,11 @@ namespace faiss {
  * RangeSearchResult
  ***********************************************************************/
 
-RangeSearchResult::RangeSearchResult(size_t nq, bool alloc_lims) : nq(nq) {
+RangeSearchResult::RangeSearchResult(size_t nq_in, bool alloc_lims)
+        : nq(nq_in) {
     if (alloc_lims) {
-        lims = new size_t[nq + 1];
-        memset(lims, 0, sizeof(*lims) * (nq + 1));
+        lims = new size_t[nq_in + 1];
+        memset(lims, 0, sizeof(*lims) * (nq_in + 1));
     } else {
         lims = nullptr;
     }
@@ -39,7 +40,7 @@ void RangeSearchResult::do_allocation() {
     // simultaneously
     FAISS_THROW_IF_NOT(labels == nullptr && distances == nullptr);
     size_t ofs = 0;
-    for (int i = 0; i < nq; i++) {
+    for (size_t i = 0; i < nq; i++) {
         size_t n = lims[i];
         lims[i] = ofs;
         ofs += n;
@@ -59,12 +60,12 @@ RangeSearchResult::~RangeSearchResult() {
  * BufferList
  ***********************************************************************/
 
-BufferList::BufferList(size_t buffer_size) : buffer_size(buffer_size) {
-    wp = buffer_size;
+BufferList::BufferList(size_t buffer_size_in) : buffer_size(buffer_size_in) {
+    wp = buffer_size_in;
 }
 
 BufferList::~BufferList() {
-    for (int i = 0; i < buffers.size(); i++) {
+    for (size_t i = 0; i < buffers.size(); i++) {
         delete[] buffers[i].ids;
         delete[] buffers[i].dis;
     }
@@ -140,7 +141,7 @@ void RangeSearchPartialResult::finalize() {
 
 /// called by range_search before do_allocation
 void RangeSearchPartialResult::set_lims() {
-    for (int i = 0; i < queries.size(); i++) {
+    for (size_t i = 0; i < queries.size(); i++) {
         RangeQueryResult& qres = queries[i];
         res->lims[qres.qno] = qres.nres;
     }
@@ -149,7 +150,7 @@ void RangeSearchPartialResult::set_lims() {
 /// called by range_search after do_allocation
 void RangeSearchPartialResult::copy_result(bool incremental) {
     size_t ofs = 0;
-    for (int i = 0; i < queries.size(); i++) {
+    for (size_t i = 0; i < queries.size(); i++) {
         RangeQueryResult& qres = queries[i];
 
         copy_range(

data/vendor/faiss/faiss/impl/AuxIndexStructures.h

@@ -18,6 +18,8 @@
 #include <mutex>
 #include <vector>
 
+#include <faiss/impl/InvertedListScannerStats.h>
+
 #include <faiss/MetricType.h>
 #include <faiss/impl/platform_macros.h>
 
@@ -92,6 +94,7 @@ struct RangeQueryResult {
     idx_t qno;   //< id of the query
     size_t nres; //< nb of results for this query
     RangeSearchPartialResult* pres;
+    InvertedListScannerStats stats;
 
     /// called by search function to report a new result
     void add(float dis, idx_t id);
@@ -169,34 +172,6 @@ struct TimeoutCallback : InterruptCallback {
     static void reset(double timeout_in_seconds);
 };
 
-/// set implementation optimized for fast access.
-struct VisitedTable {
-    std::vector<uint8_t> visited;
-    uint8_t visno;
-
-    explicit VisitedTable(int size) : visited(size), visno(1) {}
-
-    /// set flag #no to true
-    void set(int no) {
-        visited[no] = visno;
-    }
-
-    /// get flag #no
-    bool get(int no) const {
-        return visited[no] == visno;
-    }
-
-    /// reset all flags to false
-    void advance() {
-        visno++;
-        if (visno == 250) {
-            // 250 rather than 255 because sometimes we use visno and visno+1
-            memset(visited.data(), 0, sizeof(visited[0]) * visited.size());
-            visno = 1;
-        }
-    }
-};
-
 } // namespace faiss
 
 #endif

data/vendor/faiss/faiss/impl/ClusteringHelpers.cpp

@@ -0,0 +1,244 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <faiss/impl/ClusteringHelpers.h>
+
+#include <cassert>
+#include <chrono>
+#include <cinttypes>
+#include <cstdint>
+#include <cstdio>
+#include <cstring>
+#include <limits>
+#include <vector>
+
+#include <omp.h>
+
+#include <faiss/Index.h>
+#include <faiss/impl/FaissAssert.h>
+#include <faiss/utils/random.h>
+
+namespace faiss {
+namespace detail {
+
+uint64_t get_actual_rng_seed(const int seed) {
+    return (seed >= 0)
+            ? seed
+            : static_cast<uint64_t>(std::chrono::high_resolution_clock::now()
+                                            .time_since_epoch()
+                                            .count());
+}
+
+idx_t subsample_training_set(
+        const Clustering& clus,
+        idx_t nx,
+        const uint8_t* x,
+        size_t line_size,
+        const float* weights,
+        uint8_t** x_out,
+        float** weights_out) {
+    FAISS_THROW_IF_NOT(clus.k > 0 && clus.max_points_per_centroid > 0);
+    if (clus.verbose) {
+        printf("Sampling a subset of %zd / %" PRId64 " for training\n",
+               clus.k * clus.max_points_per_centroid,
+               nx);
+    }
+
+    const uint64_t actual_seed = get_actual_rng_seed(clus.seed);
+
+    std::vector<idx_t> perm;
+    if (clus.use_faster_subsampling) {
+        SplitMix64RandomGenerator rng(actual_seed);
+
+        const idx_t new_nx = clus.k * clus.max_points_per_centroid;
+        perm.resize(new_nx);
+        assert(!perm.empty());
+        for (idx_t i = 0; i < new_nx; i++) {
+            perm[i] = rng.rand_int64() % nx;
+        }
+    } else {
+        FAISS_THROW_IF_NOT_FMT(
+                nx <= static_cast<idx_t>(std::numeric_limits<int>::max()),
+                "Dataset too large (%" PRId64
+                ") for standard subsampling; "
+                "set use_faster_subsampling=true",
+                nx);
+        std::vector<int> int_perm(nx);
+        rand_perm(int_perm.data(), nx, actual_seed);
+        perm.assign(int_perm.begin(), int_perm.end());
+    }
+
+    nx = clus.k * clus.max_points_per_centroid;
+    FAISS_THROW_IF_NOT_FMT(
+            perm.size() >= static_cast<size_t>(nx),
+            "subsample_training_set: perm size %zu < required nx %" PRId64,
+            perm.size(),
+            nx);
+    assert(!perm.empty());
+
+    uint8_t* x_new = new uint8_t[nx * line_size];
+    *x_out = x_new;
+
+    for (idx_t i = 0; i < nx; i++) {
+        memcpy(x_new + i * line_size, x + perm[i] * line_size, line_size);
+    }
+    if (weights) {
+        float* weights_new = new float[nx];
+        for (idx_t i = 0; i < nx; i++) {
+            weights_new[i] = weights[perm[i]];
+        }
+        *weights_out = weights_new;
+    } else {
+        *weights_out = nullptr;
+    }
+    return nx;
+}
+
+void compute_centroids(
+        size_t d,
+        size_t k,
+        size_t n,
+        size_t k_frozen,
+        const uint8_t* x,
+        const Index* codec,
+        const int64_t* assign,
+        const float* weights,
+        float* hassign,
+        float* centroids) {
+    k -= k_frozen;
+    centroids += k_frozen * d;
+
+    memset(centroids, 0, sizeof(*centroids) * d * k);
+
+    size_t line_size = codec ? codec->sa_code_size() : d * sizeof(float);
+
+#pragma omp parallel
+    {
+        int nt = omp_get_num_threads();
+        int rank = omp_get_thread_num();
+
+        // this thread is taking care of centroids c0:c1
+        size_t c0 = (k * rank) / nt;
+        size_t c1 = (k * (rank + 1)) / nt;
+        std::vector<float> decode_buffer(d);
+
+        for (size_t i = 0; i < n; i++) {
+            int64_t ci = assign[i];
+            FAISS_THROW_IF_NOT_MSG(
+                    ci >= 0 && ci < k + k_frozen, "invalid cluster assignment");
+            ci -= k_frozen;
+            if (ci >= static_cast<int64_t>(c0) &&
+                ci < static_cast<int64_t>(c1)) {
+                float* c = centroids + ci * d;
+                const float* xi;
+                if (!codec) {
+                    xi = reinterpret_cast<const float*>(x + i * line_size);
+                } else {
+                    float* xif = decode_buffer.data();
+                    codec->sa_decode(1, x + i * line_size, xif);
+                    xi = xif;
+                }
+                if (weights) {
+                    float w = weights[i];
+                    hassign[ci] += w;
+                    for (size_t j = 0; j < d; j++) {
+                        c[j] += xi[j] * w;
+                    }
+                } else {
+                    hassign[ci] += 1.0;
+                    for (size_t j = 0; j < d; j++) {
+                        c[j] += xi[j];
+                    }
+                }
+            }
+        }
+    }
+
+#pragma omp parallel for
+    for (idx_t ci = 0; ci < static_cast<idx_t>(k); ci++) {
+        if (hassign[ci] == 0) {
+            continue;
+        }
+        float norm = 1 / hassign[ci];
+        float* c = centroids + ci * d;
+        for (size_t j = 0; j < d; j++) {
+            c[j] *= norm;
+        }
+    }
+}
+
+// a bit above machine epsilon for float16
+static constexpr float EPS = 1.f / 1024.f;
+
+int split_clusters(
+        size_t d,
+        size_t k,
+        size_t n,
+        size_t k_frozen,
+        float* hassign,
+        float* centroids) {
+    k -= k_frozen;
+    centroids += k_frozen * d;
+    FAISS_THROW_IF_NOT_MSG(
+            n > k,
+            "split_clusters: n must exceed k to find a non-empty donor centroid");
+
+    size_t nsplit = 0;
+    RandomGenerator rng(1234);
+    for (size_t ci = 0; ci < k; ci++) {
+        if (hassign[ci] == 0) {
+            // Probabilistic donor pick weighted by hassign; deterministic
+            // fallback to the largest cluster if too many iterations pass.
+            size_t cj;
+            size_t max_tries = 10 * k;
+            size_t n_tries = 0;
+            bool found = false;
+            for (cj = 0; n_tries < max_tries; cj = (cj + 1) % k) {
+                float p = (hassign[cj] - 1.0) / (float)(n - k);
+                float r = rng.rand_float();
+                if (r < p) {
+                    found = true;
+                    break;
+                }
+                n_tries++;
+            }
+            if (!found) {
+                // Deterministic fallback: split the largest cluster.
+                cj = 0;
+                for (size_t j = 1; j < k; j++) {
+                    if (hassign[j] > hassign[cj]) {
+                        cj = j;
+                    }
+                }
+            }
+            memcpy(centroids + ci * d,
+                   centroids + cj * d,
+                   sizeof(*centroids) * d);
+
+            /* small symmetric perturbation */
+            for (size_t j = 0; j < d; j++) {
+                if (j % 2 == 0) {
+                    centroids[ci * d + j] *= 1 + EPS;
+                    centroids[cj * d + j] *= 1 - EPS;
+                } else {
+                    centroids[ci * d + j] *= 1 - EPS;
+                    centroids[cj * d + j] *= 1 + EPS;
+                }
+            }
+
+            /* assume even split of the cluster */
+            hassign[ci] = hassign[cj] / 2;
+            hassign[cj] -= hassign[ci];
+            nsplit++;
+        }
+    }
+
+    return static_cast<int>(nsplit);
+}
+
+} // namespace detail
+} // namespace faiss

data/vendor/faiss/faiss/impl/ClusteringHelpers.h

@@ -0,0 +1,94 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+
+#include <faiss/Clustering.h>
+#include <faiss/Index.h>
+
+namespace faiss {
+namespace detail {
+
+/** Resolve the actual RNG seed for clustering helpers.
+ *
+ * If `seed >= 0`, returns `seed`. Otherwise returns a high-resolution
+ * timestamp so that callers get a non-deterministic seed.
+ *
+ * @param seed  user-provided seed; negative values request a time-based seed
+ * @return      the resolved seed
+ */
+uint64_t get_actual_rng_seed(const int seed);
+
+/** Subsample a training set down to `clus.k * clus.max_points_per_centroid`
+ * rows.
+ *
+ * Allocates `*x_out` (and `*weights_out` when `weights` is non-null) with
+ * `new[]`; ownership is transferred to the caller.
+ *
+ * @param clus        clustering parameters (reads `k`,
+ * `max_points_per_centroid`, `use_faster_subsampling`, `seed`, `verbose`)
+ * @param nx          number of input training rows
+ * @param x           input training data, row-major, `nx * line_size` bytes
+ * @param line_size   bytes per training row
+ * @param weights     optional per-row weights (length `nx`), or null
+ * @param x_out       output: newly allocated subsampled rows
+ * @param weights_out output: newly allocated subsampled weights, or null
+ * @return            number of rows in the subsampled set
+ */
+idx_t subsample_training_set(
+        const Clustering& clus,
+        idx_t nx,
+        const uint8_t* x,
+        size_t line_size,
+        const float* weights,
+        uint8_t** x_out,
+        float** weights_out);
+
+/** compute centroids as (weighted) sum of training points
+ *
+ * @param x            training vectors, size n * code_size (from codec)
+ * @param codec        how to decode the vectors (if NULL then cast to float*)
+ * @param weights      per-training vector weight, size n (or NULL)
+ * @param assign       nearest centroid for each training vector, size n
+ * @param k_frozen     do not update the k_frozen first centroids
+ * @param centroids    centroid vectors (output only), size k * d
+ * @param hassign      histogram of assignments per centroid (size k),
+ *                     should be 0 on input
+ *
+ */
+void compute_centroids(
+        size_t d,
+        size_t k,
+        size_t n,
+        size_t k_frozen,
+        const uint8_t* x,
+        const Index* codec,
+        const int64_t* assign,
+        const float* weights,
+        float* hassign,
+        float* centroids);
+
+/** Handle empty clusters by splitting larger ones.
+ *
+ * It works by slightly changing the centroids to make 2 clusters from
+ * a single one. Takes the same arguments as compute_centroids.
+ *
+ * @return           nb of splitting operations (larger is worse)
+ */
+int split_clusters(
+        size_t d,
+        size_t k,
+        size_t n,
+        size_t k_frozen,
+        float* hassign,
+        float* centroids);
+
+} // namespace detail
+} // namespace faiss