faiss 0.6.4 → 0.6.5

data/vendor/faiss/faiss/gpu_metal/MetalDistance.h

@@ -83,5 +83,97 @@ bool runMetalIVFFlatFullSearch(
         bool centroidsAreFP16 = false,
         bool waitForCompletion = true);
 
+/// Full IVF-PQ search on the GPU: builds per-(query, probe) PQ lookup tables,
+/// scans the assigned inverted lists with 8-bit ADC, and merges the per-list
+/// top-k into the final top-k. Coarse quantization is performed on the CPU and
+/// the assignments are passed in via @p coarseAssign. @p coarseDist carries
+/// the per-(query, probe) coarse scores, added for inner product (the IP
+/// tables hold <q_m, pq_m> only); unused for L2.
+bool runMetalIVFPQFullSearch(
+        id<MTLDevice> device,
+        id<MTLCommandQueue> queue,
+        id<MTLBuffer> queries,
+        id<MTLBuffer> coarseAssign,
+        id<MTLBuffer> coarseDist,
+        id<MTLBuffer> coarseCentroids,
+        id<MTLBuffer> pqCentroids,
+        id<MTLBuffer> lookupTable,
+        id<MTLBuffer> codes,
+        id<MTLBuffer> ids,
+        id<MTLBuffer> listOffset,
+        id<MTLBuffer> listLength,
+        int nq,
+        int d,
+        int M,
+        int k,
+        int nprobe,
+        int nlist,
+        int avgListLen,
+        bool lookupFp16,
+        bool isL2,
+        id<MTLBuffer> outDistances,
+        id<MTLBuffer> outIndices,
+        id<MTLBuffer> perListDistBuf,
+        id<MTLBuffer> perListIdxBuf,
+        bool waitForCompletion = true);
+
+/// Build the query-independent precomputed term of the IVFPQ L2 distance
+/// decomposition on the GPU: term2[l][m][c] = ||pq[m][c]||^2 +
+/// 2 <coarse_l[m], pq[m][c]>, laid out as nlist * M * 256 floats. Run once
+/// per trained index; consumed by runMetalIVFPQPrecompSearch.
+bool runMetalIVFPQPrecomputeTerm2(
+        id<MTLDevice> device,
+        id<MTLCommandQueue> queue,
+        id<MTLBuffer> coarseCentroids,
+        id<MTLBuffer> pqCentroids,
+        id<MTLBuffer> outTerm2,
+        int nlist,
+        int d,
+        int M);
+
+/// IVF-PQ search using the precomputed-table decomposition: a per-query
+/// M*256 term is built once per batch (instead of a full residual LUT per
+/// (query, probe)), combined with @p term2 in threadgroup memory by the scan
+/// kernel, and the coarse distance is added as the constant ||x - c||^2 term.
+/// The scan keeps an exact running top-k over arbitrarily long lists and the
+/// merge runs in rounds, so there is no list-length or nprobe*k cap; the
+/// remaining requirements are k <= 512, M <= 64 and d/M <= 256.
+/// onTheFly evaluates encountered codes directly (d <= 256), using the
+/// supplied coarse centroids; shortLists requires every list to have <= 128
+/// entries and selects a smaller threadgroup and scratch allocation.
+/// @p term2 may be nil for inner product (qterm is the whole table);
+/// @p useDis0 adds the coarse distance (required for L2, by_residual for IP).
+bool runMetalIVFPQPrecompSearch(
+        id<MTLDevice> device,
+        id<MTLCommandQueue> queue,
+        id<MTLBuffer> queries,
+        id<MTLBuffer> coarseAssign,
+        id<MTLBuffer> coarseDist,
+        id<MTLBuffer> term2,
+        id<MTLBuffer> qtermScratch,
+        id<MTLBuffer> pqCentroids,
+        id<MTLBuffer> codes,
+        id<MTLBuffer> ids,
+        id<MTLBuffer> listOffset,
+        id<MTLBuffer> listLength,
+        int nq,
+        int d,
+        int M,
+        int k,
+        int nprobe,
+        bool isL2,
+        bool useDis0,
+        id<MTLBuffer> outDistances,
+        id<MTLBuffer> outIndices,
+        id<MTLBuffer> perListDistBuf,
+        id<MTLBuffer> perListIdxBuf,
+        id<MTLBuffer> mergeScratchDistBuf,
+        id<MTLBuffer> mergeScratchIdxBuf,
+        bool waitForCompletion = true,
+        id<MTLBuffer> coarseCentroids = nil,
+        bool onTheFly = false,
+        bool shortLists = false,
+        bool compactMerge = true);
+
 } // namespace gpu_metal
 } // namespace faiss

data/vendor/faiss/faiss/gpu_metal/MetalIndex.h

@@ -24,6 +24,10 @@ struct MetalIndexConfig {
 
     bool useFloat16CoarseQuantizer = false;
 
+    /// Build the IVF-PQ distance lookup tables in half precision. Halves LUT
+    /// bandwidth during the GPU scan at a small accuracy cost.
+    bool useFloat16 = false;
+
     faiss::gpu::IndicesOptions indicesOptions = faiss::gpu::INDICES_64_BIT;
 
     bool interleavedLayout = true;

data/vendor/faiss/faiss/gpu_metal/MetalIndexBinaryFlat.h

@@ -0,0 +1,69 @@
+// @lint-ignore-every LICENSELINT
+/**
+ * Copyright (c) Meta Platforms, Inc. and its 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 <faiss/IndexBinaryFlat.h>
+#include <faiss/gpu_metal/MetalResources.h>
+#include <memory>
+
+namespace faiss {
+namespace gpu_metal {
+
+/// Exhaustive Hamming search over GPU-resident binary vectors.
+/// Every positive dimension divisible by 8 is supported. GPU search is exact
+/// for k <= 512, with ties ordered by increasing label. Larger k, selectors,
+/// and counts exceeding uint32_t use the CPU copy. GPU failures throw rather
+/// than silently falling back. Range search and remove_ids are not supported.
+/// As with other Metal indexes, callers must serialize mutation and search.
+class MetalIndexBinaryFlat : public faiss::IndexBinary {
+   public:
+    static constexpr int kMaxGpuK = 512;
+
+    MetalIndexBinaryFlat(std::shared_ptr<MetalResources> resources, int d);
+    MetalIndexBinaryFlat(
+            std::shared_ptr<MetalResources> resources,
+            const faiss::IndexBinaryFlat* index);
+    ~MetalIndexBinaryFlat() override;
+
+    MetalIndexBinaryFlat(const MetalIndexBinaryFlat&) = delete;
+    MetalIndexBinaryFlat& operator=(const MetalIndexBinaryFlat&) = delete;
+
+    void add(idx_t n, const uint8_t* x) override;
+    void reset() override;
+    void search(
+            idx_t n,
+            const uint8_t* x,
+            idx_t k,
+            int32_t* distances,
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
+    void reconstruct(idx_t key, uint8_t* recons) const override;
+
+    void copyFrom(const faiss::IndexBinaryFlat* index);
+    void copyTo(faiss::IndexBinaryFlat* index) const;
+
+    std::shared_ptr<MetalResources> getResources() const {
+        return resources_;
+    }
+
+   private:
+    std::shared_ptr<MetalResources> resources_;
+    faiss::IndexBinaryFlat cpuIndex_;
+    id<MTLBuffer> vectorsBuf_ = nil;
+    // Search is synchronous and callers serialize access, so scratch buffers
+    // can be reused after waitUntilCompleted. ARC releases them on destruction.
+    mutable id<MTLBuffer> queryBuf_ = nil;
+    mutable id<MTLBuffer> distanceBuf_ = nil;
+    mutable id<MTLBuffer> labelBuf_ = nil;
+
+    id<MTLBuffer> upload(const uint8_t* data, size_t bytes) const;
+};
+
+} // namespace gpu_metal
+} // namespace faiss

data/vendor/faiss/faiss/gpu_metal/MetalIndexIVFPQ.h

@@ -75,13 +75,51 @@ class MetalIndexIVFPQ : public MetalIndex {
     std::unique_ptr<faiss::IndexIVFPQ> cpuIndex_;
     std::unique_ptr<MetalIVFPQImpl> gpuIvf_;
 
+    // Reusable GPU scratch buffers for the search path, grown on demand.
+    mutable id<MTLBuffer> searchQueriesBuf_ = nil;
+    mutable id<MTLBuffer> searchCoarseBuf_ = nil;
+    mutable id<MTLBuffer> searchOutDistBuf_ = nil;
+    mutable id<MTLBuffer> searchOutIdxBuf_ = nil;
+    mutable id<MTLBuffer> searchPerListDistBuf_ = nil;
+    mutable id<MTLBuffer> searchPerListIdxBuf_ = nil;
+    mutable id<MTLBuffer> lookupTableBuf_ = nil;
+    mutable id<MTLBuffer> lookupTableHalfBuf_ = nil;
+    mutable id<MTLBuffer> searchCoarseDistBuf_ = nil;
+    mutable id<MTLBuffer> searchQTermBuf_ = nil;
+    mutable id<MTLBuffer> searchMergeDistBuf_ = nil;
+    mutable id<MTLBuffer> searchMergeIdxBuf_ = nil;
+    mutable size_t searchQueriesCap_ = 0;
+    mutable size_t searchCoarseCap_ = 0;
+    mutable size_t searchOutDistCap_ = 0;
+    mutable size_t searchOutIdxCap_ = 0;
+    mutable size_t searchPerListDistCap_ = 0;
+    mutable size_t searchPerListIdxCap_ = 0;
+    mutable size_t lookupTableCap_ = 0;
+    mutable size_t lookupTableHalfCap_ = 0;
+    mutable size_t searchCoarseDistCap_ = 0;
+    mutable size_t searchQTermCap_ = 0;
+    mutable size_t searchMergeDistCap_ = 0;
+    mutable size_t searchMergeIdxCap_ = 0;
+
+    // Coarse centroids mirrored on the GPU for the L2 lookup-table build.
+    mutable id<MTLBuffer> centroidBuf_ = nil;
+
+    // Query-independent term of the L2 precomputed-table decomposition
+    // (nlist * M * 256 floats), built once per trained index on the GPU.
+    mutable id<MTLBuffer> term2Buf_ = nil;
+
     void verifyPQSettings_() const;
+    void precomputeTerm2_() const;
 
     void encodeResidualAndAppend_(
             idx_t n,
             const float* x,
             const idx_t* list_nos,
             const idx_t* xids);
+
+    void ensureSearchBuf_(id<MTLBuffer>& buf, size_t& cap, size_t needed) const;
+    void uploadCentroids_() const;
+    void uploadPQCentroids_() const;
 };
 
 } // namespace gpu_metal

data/vendor/faiss/faiss/gpu_metal/MetalKernels.h

@@ -91,6 +91,112 @@ class MetalKernels {
             id<MTLBuffer> paramsBuf,
             int nq);
 
+    // ---- IVF-PQ ----
+
+    void encodeIVFPQBuildLookupTables(
+            id<MTLComputeCommandEncoder> enc,
+            bool isL2,
+            bool outFp16,
+            id<MTLBuffer> queries,
+            id<MTLBuffer> coarseAssign,
+            id<MTLBuffer> coarseCentroids,
+            id<MTLBuffer> pqCentroids,
+            id<MTLBuffer> outLookup,
+            int nq,
+            int d,
+            int M,
+            int nprobe);
+
+    void encodeIVFPQScanList(
+            id<MTLComputeCommandEncoder> enc,
+            bool useSmall,
+            bool useFp16Lookup,
+            id<MTLBuffer> lookupTable,
+            id<MTLBuffer> codes,
+            id<MTLBuffer> ids,
+            id<MTLBuffer> listOffset,
+            id<MTLBuffer> listLength,
+            id<MTLBuffer> coarseAssign,
+            id<MTLBuffer> perListDist,
+            id<MTLBuffer> perListIdx,
+            id<MTLBuffer> paramsBuf,
+            id<MTLBuffer> coarseDist,
+            int nq,
+            int nprobe);
+
+    // ---- IVF-PQ precomputed-table path ----
+
+    void encodeIVFPQPrecomputeTerm2(
+            id<MTLComputeCommandEncoder> enc,
+            id<MTLBuffer> coarseCentroids,
+            id<MTLBuffer> pqCentroids,
+            id<MTLBuffer> outTerm2,
+            int nlist,
+            int d,
+            int M);
+
+    void encodeIVFPQBuildQueryTerm(
+            id<MTLComputeCommandEncoder> enc,
+            id<MTLBuffer> queries,
+            id<MTLBuffer> pqCentroids,
+            id<MTLBuffer> outQTerm,
+            int nq,
+            int d,
+            int M,
+            bool isL2);
+
+    void encodeIVFPQScanListPrecomp(
+            id<MTLComputeCommandEncoder> enc,
+            id<MTLBuffer> term2, // nil when useTerm2 is false (IP path)
+            id<MTLBuffer> qterm,
+            id<MTLBuffer> coarseDist,
+            id<MTLBuffer> codes,
+            id<MTLBuffer> ids,
+            id<MTLBuffer> listOffset,
+            id<MTLBuffer> listLength,
+            id<MTLBuffer> coarseAssign,
+            id<MTLBuffer> perListDist,
+            id<MTLBuffer> perListIdx,
+            int nq,
+            int M,
+            int k,
+            int nprobe,
+            bool wantMin,
+            bool useTerm2,
+            bool useDis0,
+            id<MTLBuffer> queries,
+            id<MTLBuffer> coarseCentroids,
+            id<MTLBuffer> pqCentroids,
+            int d,
+            bool onTheFly,
+            bool shortLists);
+
+    void encodeIVFMergeListsGrouped(
+            id<MTLComputeCommandEncoder> enc,
+            id<MTLBuffer> inDist,
+            id<MTLBuffer> inIdx,
+            id<MTLBuffer> outDist,
+            id<MTLBuffer> outIdx,
+            int nq,
+            int numLists,
+            int groupSize,
+            int k,
+            bool wantMin,
+            bool compact);
+
+    // Exact Hamming top-k from the precompiled MetalDistance.metallib.
+    // Requires positive uint32-sized counts and 1 <= k <= 512.
+    void encodeHammingDistanceTopK(
+            id<MTLComputeCommandEncoder> enc,
+            id<MTLBuffer> queries,
+            id<MTLBuffer> database,
+            id<MTLBuffer> outDist,
+            id<MTLBuffer> outIdx,
+            uint32_t nq,
+            uint32_t nb,
+            uint32_t codeSize,
+            uint32_t k);
+
     static int selectTopKVariantIndex(int k);
 
    private:

data/vendor/faiss/faiss/gpu_metal/MetalPythonBridge.h

@@ -12,6 +12,7 @@
 #pragma once
 
 #include <faiss/Index.h>
+#include <faiss/IndexBinary.h>
 
 namespace faiss {
 namespace gpu_metal {
@@ -41,5 +42,14 @@ faiss::Index* index_cpu_to_gpu(
 /// Copy Metal index back to CPU. Caller owns returned index.
 faiss::Index* index_gpu_to_cpu(const faiss::Index* index);
 
+/// Clone CPU IndexBinaryFlat to Metal. Caller owns the returned index.
+faiss::IndexBinary* index_binary_cpu_to_gpu(
+        StandardMetalResourcesHolder* res,
+        int device,
+        const faiss::IndexBinary* index);
+
+/// Clone MetalIndexBinaryFlat back to CPU. Caller owns the returned index.
+faiss::IndexBinary* index_binary_gpu_to_cpu(const faiss::IndexBinary* index);
+
 } // namespace gpu_metal
 } // namespace faiss

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

@@ -246,6 +246,8 @@ void TimeoutCallback::set_timeout(double timeout_in_seconds) {
 }
 
 bool TimeoutCallback::want_interrupt() {
+    static std::mutex mutex;
+    std::lock_guard<std::mutex> guard(mutex);
     if (timeout == 0) {
         return false;
     }

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

@@ -146,12 +146,7 @@ struct FAISS_API InterruptCallback {
 
     static void clear_instance();
 
-    /** check if:
-     * - an interrupt callback is set
-     * - the callback returns true
-     * if this is the case, then throw an exception. Should not be called
-     * from multiple threads.
-     */
+    /// Concurrent calls require a thread-safe want_interrupt().
     static void check();
 
     /// same as check() but return true if is interrupted instead of

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

@@ -7,6 +7,8 @@
 
 #include <faiss/impl/ClusteringHelpers.h>
 
+#include <algorithm>
+#include <atomic>
 #include <cassert>
 #include <chrono>
 #include <cinttypes>
@@ -20,11 +22,88 @@
 
 #include <faiss/Index.h>
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/simd_dispatch.h>
 #include <faiss/utils/random.h>
+#include <faiss/utils/simd_impl/fp16_kernels.h>
 
 namespace faiss {
 namespace detail {
 
+namespace {
+
+#if defined(_OPENMP) && !defined(FAISS_USE_OMP_MOCK)
+/// Below this many values, the fp16 finiteness check stays on one thread.
+constexpr size_t fp16_parallel_threshold = size_t(1) << 16;
+#endif
+
+/** Shared body of compute_centroids and compute_centroids_fp16.
+ *
+ * `add_row(i, w, c, scratch)` adds training vector i to the fp32 vector c,
+ * scaled by *w when w is not null. `scratch` is a per-thread buffer of d
+ * floats.
+ */
+template <class AddRow>
+void compute_centroids_impl(
+        size_t d,
+        size_t k,
+        size_t n,
+        size_t k_frozen,
+        const int64_t* assign,
+        const float* weights,
+        float* hassign,
+        float* centroids,
+        const AddRow& add_row) {
+    k -= k_frozen;
+    centroids += k_frozen * d;
+
+    memset(centroids, 0, sizeof(*centroids) * d * k);
+
+    std::atomic<bool> invalid_assignment{false};
+    const int64_t num_centroids = static_cast<int64_t>(k + k_frozen);
+#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> scratch(d);
+
+        for (size_t i = 0; i < n; i++) {
+            int64_t ci = assign[i];
+            if (ci < 0 || ci >= num_centroids) {
+                invalid_assignment.store(true, std::memory_order_relaxed);
+                continue;
+            }
+            ci -= k_frozen;
+            if (ci >= static_cast<int64_t>(c0) &&
+                ci < static_cast<int64_t>(c1)) {
+                const float* w = weights ? weights + i : nullptr;
+                hassign[ci] += w ? *w : 1.0f;
+                add_row(i, w, centroids + ci * d, scratch.data());
+            }
+        }
+    }
+    FAISS_THROW_IF_MSG(
+            invalid_assignment.load(std::memory_order_relaxed),
+            "invalid cluster assignment");
+
+#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;
+        }
+    }
+}
+
+} // namespace
+
 uint64_t get_actual_rng_seed(const int seed) {
     return (seed >= 0)
             ? seed
@@ -109,66 +188,76 @@ void compute_centroids(
         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 size_t line_size = codec ? codec->sa_code_size() : d * sizeof(float);
+    compute_centroids_impl(
+            d,
+            k,
+            n,
+            k_frozen,
+            assign,
+            weights,
+            hassign,
+            centroids,
+            [&](size_t i, const float* w, float* c, float* scratch) {
                 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;
+                    codec->sa_decode(1, x + i * line_size, scratch);
+                    xi = scratch;
                 }
-                if (weights) {
-                    float w = weights[i];
-                    hassign[ci] += w;
+                if (w) {
+                    const float wi = *w;
                     for (size_t j = 0; j < d; j++) {
-                        c[j] += xi[j] * w;
+                        c[j] += xi[j] * wi;
                     }
                 } 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;
+void compute_centroids_fp16(
+        size_t d,
+        size_t k,
+        size_t n,
+        size_t k_frozen,
+        const uint16_t* x,
+        const int64_t* assign,
+        const float* weights,
+        float* hassign,
+        float* centroids) {
+    with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_BASE>([&]<SIMDLevel SL>() {
+        compute_centroids_impl(
+                d,
+                k,
+                n,
+                k_frozen,
+                assign,
+                weights,
+                hassign,
+                centroids,
+                [&](size_t i, const float* w, float* c, float* /*scratch*/) {
+                    // x * 1.0f is exact, so unweighted sums match the
+                    // fp32 path.
+                    fp16_madd<SL>(d, x + i * d, w ? *w : 1.0f, c);
+                });
+    });
+}
+
+bool fp16_all_finite(size_t n, const uint16_t* x) {
+    // binary16 NaNs and infinities are the values with all exponent bits set
+    constexpr uint16_t exponent_mask = 0x7c00;
+    std::atomic<bool> all_finite{true};
+#pragma omp parallel for if (n > fp16_parallel_threshold)
+    for (int64_t i = 0; i < static_cast<int64_t>(n); i++) {
+        if ((x[i] & exponent_mask) == exponent_mask) {
+            all_finite.store(false, std::memory_order_relaxed);
         }
     }
+    return all_finite.load(std::memory_order_relaxed);
 }
 
 // a bit above machine epsilon for float16

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

@@ -51,6 +51,15 @@ idx_t subsample_training_set(
         uint8_t** x_out,
         float** weights_out);
 
+/** Initialize non-random centroids directly from packed IEEE binary16 rows. */
+void init_centroids_fp16(
+        const ClusteringInitialization& initializer,
+        size_t n,
+        const uint16_t* x,
+        float* centroids,
+        size_t n_existing_centroids = 0,
+        const float* existing_centroids = nullptr);
+
 /** compute centroids as (weighted) sum of training points
  *
  * @param x            training vectors, size n * code_size (from codec)
@@ -75,6 +84,33 @@ void compute_centroids(
         float* hassign,
         float* centroids);
 
+/** Same as compute_centroids for training vectors stored as packed IEEE
+ * binary16 values.
+ *
+ * Rows are widened in registers and accumulated directly into the fp32
+ * centroids: there is no per-row decode call and no decode buffer.
+ *
+ * @param x            training vectors, size n * d
+ */
+void compute_centroids_fp16(
+        size_t d,
+        size_t k,
+        size_t n,
+        size_t k_frozen,
+        const uint16_t* x,
+        const int64_t* assign,
+        const float* weights,
+        float* hassign,
+        float* centroids);
+
+/** Check packed IEEE binary16 values for NaNs and infinities.
+ *
+ * The check runs on the bit patterns, without widening to fp32.
+ *
+ * @return true iff none of the `n` values is a NaN or an infinity
+ */
+bool fp16_all_finite(size_t n, const uint16_t* x);
+
 /** Handle empty clusters by splitting larger ones.
  *
  * It works by slightly changing the centroids to make 2 clusters from