faiss 0.5.3 → 0.6.1

data/vendor/faiss/faiss/svs/IndexSVSIVFLeanVec.cpp

@@ -0,0 +1,225 @@
+/*
+ * Portions 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.
+ */
+
+/*
+ * Portions Copyright 2026 Intel Corporation
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <faiss/svs/IndexSVSFaissUtils.h>
+#include <faiss/svs/IndexSVSIVFLeanVec.h>
+
+#include <svs/runtime/dynamic_ivf_index.h>
+#include <svs/runtime/ivf_index.h>
+#include <svs/runtime/training.h>
+
+#include <memory>
+#include <numeric>
+
+namespace faiss {
+
+IndexSVSIVFLeanVec::IndexSVSIVFLeanVec() : IndexSVSIVF() {
+    is_trained = false;
+    storage_kind = SVSStorageKind::SVS_LeanVec4x4;
+}
+
+IndexSVSIVFLeanVec::IndexSVSIVFLeanVec(
+        idx_t d,
+        size_t nlist,
+        MetricType metric,
+        size_t leanvec_dims,
+        SVSStorageKind storage_kind,
+        bool is_static)
+        : IndexSVSIVF(d, nlist, metric, storage_kind, is_static) {
+    is_trained = false;
+    leanvec_d = leanvec_dims == 0 ? d / 2 : leanvec_dims;
+}
+
+IndexSVSIVFLeanVec::~IndexSVSIVFLeanVec() {
+    if (training_data) {
+        auto status = svs_runtime::LeanVecTrainingData::destroy(training_data);
+        FAISS_ASSERT(status.ok());
+        training_data = nullptr;
+    }
+    // Base class destructor handles impl cleanup
+}
+
+void IndexSVSIVFLeanVec::train(idx_t n, const float* x) {
+    train_with_queries(n, x, 0, nullptr);
+}
+
+void IndexSVSIVFLeanVec::train_with_queries(
+        idx_t n,
+        const float* x,
+        idx_t n_train_q,
+        const float* xq_train) {
+    FAISS_THROW_IF_MSG(
+            training_data || impl, "Index already trained or contains data.");
+
+    FAISS_THROW_IF_NOT_MSG(
+            IndexSVSIVF::is_lvq_leanvec_enabled(),
+            "LVQ/LeanVec support not available on this platform or build");
+
+    // Build LeanVec training data
+    auto status = svs_runtime::LeanVecTrainingData::build(
+            &training_data, d, n, x, n_train_q, xq_train, leanvec_d);
+    if (!status.ok()) {
+        FAISS_THROW_MSG(status.message());
+    }
+    FAISS_THROW_IF_NOT_MSG(
+            training_data, "Failed to build leanvec training info.");
+
+    // Now build the IVF index with the training data
+    create_impl(n, x);
+    is_trained = true;
+}
+
+void IndexSVSIVFLeanVec::reset() {
+    if (training_data) {
+        auto status = svs_runtime::LeanVecTrainingData::destroy(training_data);
+        FAISS_ASSERT(status.ok());
+        training_data = nullptr;
+    }
+    IndexSVSIVF::reset();
+}
+
+void IndexSVSIVFLeanVec::serialize_training_data(std::ostream& out) const {
+    FAISS_THROW_IF_NOT_MSG(
+            training_data, "Cannot serialize: Training data not initialized.");
+
+    auto status = training_data->save(out);
+    if (!status.ok()) {
+        FAISS_THROW_MSG(status.message());
+    }
+}
+
+void IndexSVSIVFLeanVec::deserialize_training_data(std::istream& in) {
+    svs_runtime::LeanVecTrainingData* tdata = nullptr;
+    auto status = svs_runtime::LeanVecTrainingData::load(&tdata, in);
+    if (!status.ok()) {
+        FAISS_THROW_MSG(status.message());
+    }
+    FAISS_THROW_IF_NOT_MSG(tdata, "Failed to load leanvec training data.");
+    training_data = tdata;
+}
+
+void IndexSVSIVFLeanVec::create_impl(idx_t n, const float* x) {
+    FAISS_THROW_IF_NOT(!impl);
+    ntotal = 0;
+    auto svs_metric = to_svs_metric(metric_type);
+    auto svs_storage_kind = to_svs_storage_kind(storage_kind);
+    auto build_params = svs_runtime::IVFIndex::BuildParams{
+            .num_centroids = num_centroids,
+            .minibatch_size = minibatch_size,
+            .num_iterations = num_iterations,
+            .is_hierarchical = is_hierarchical,
+            .training_fraction = training_fraction,
+            .hierarchical_level1_clusters = hierarchical_level1_clusters,
+            .seed = seed,
+    };
+    auto search_params = svs_runtime::IVFIndex::SearchParams{
+            .n_probes = n_probes,
+            .k_reorder = k_reorder,
+    };
+
+    auto status = svs_runtime::Status_Ok;
+    if (is_static) {
+        if (training_data) {
+            status = svs_runtime::IVFIndexLeanVec::build(
+                    &impl,
+                    d,
+                    svs_metric,
+                    svs_storage_kind,
+                    static_cast<size_t>(n),
+                    x,
+                    training_data,
+                    build_params,
+                    search_params,
+                    num_threads,
+                    intra_query_threads);
+        } else {
+            status = svs_runtime::IVFIndexLeanVec::build(
+                    &impl,
+                    d,
+                    svs_metric,
+                    svs_storage_kind,
+                    static_cast<size_t>(n),
+                    x,
+                    leanvec_d,
+                    build_params,
+                    search_params,
+                    num_threads,
+                    intra_query_threads);
+        }
+    } else {
+        std::vector<size_t> labels(n);
+        std::iota(labels.begin(), labels.end(), 0);
+
+        svs_runtime::DynamicIVFIndex* dyn_impl = nullptr;
+        if (training_data) {
+            status = svs_runtime::DynamicIVFIndexLeanVec::build(
+                    &dyn_impl,
+                    d,
+                    svs_metric,
+                    svs_storage_kind,
+                    static_cast<size_t>(n),
+                    x,
+                    labels.data(),
+                    training_data,
+                    build_params,
+                    search_params,
+                    num_threads,
+                    intra_query_threads);
+        } else {
+            status = svs_runtime::DynamicIVFIndexLeanVec::build(
+                    &dyn_impl,
+                    d,
+                    svs_metric,
+                    svs_storage_kind,
+                    static_cast<size_t>(n),
+                    x,
+                    labels.data(),
+                    leanvec_d,
+                    build_params,
+                    search_params,
+                    num_threads,
+                    intra_query_threads);
+        }
+        impl = dyn_impl;
+    }
+
+    if (!status.ok()) {
+        FAISS_THROW_MSG(status.message());
+    }
+    FAISS_THROW_IF_NOT(impl);
+
+    // Force reinitialization of SVS search infrastructure after build.
+    // Same workaround as in IndexSVSIVF::create_impl — the move chain through
+    // the orchestrator wrappers can leave matmul_results_ inconsistent
+    // when threads > centroids.
+    size_t current_threads = 0;
+    impl->get_num_threads(&current_threads);
+    auto st = impl->set_num_threads(current_threads);
+    if (!st.ok()) {
+        FAISS_THROW_MSG(st.message());
+    }
+
+    ntotal = n;
+}
+
+} // namespace faiss

data/vendor/faiss/faiss/svs/IndexSVSIVFLeanVec.h

@@ -0,0 +1,71 @@
+/*
+ * Portions 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.
+ */
+
+/*
+ * Portions Copyright 2026 Intel Corporation
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <faiss/svs/IndexSVSIVF.h>
+
+namespace faiss {
+
+struct IndexSVSIVFLeanVec : IndexSVSIVF {
+    IndexSVSIVFLeanVec();
+
+    IndexSVSIVFLeanVec(
+            idx_t d,
+            size_t nlist,
+            MetricType metric = METRIC_L2,
+            size_t leanvec_dims = 0,
+            SVSStorageKind storage = SVSStorageKind::SVS_LeanVec4x4,
+            bool is_static = false);
+
+    ~IndexSVSIVFLeanVec() override;
+
+    /* Default train assumes in-distribution data */
+    void train(idx_t n, const float* x) override;
+
+    /* Generic train with out-of-distribution parameters.
+     * Out-of-distribution (OOD) means database vectors and queries _can_ be
+     * sampled from different distributions (e.g., cross-modal). More details in
+     * the original publication, arXiv:2312.16335.
+     */
+    void train_with_queries(
+            idx_t n,
+            const float* x,
+            idx_t n_train_q,
+            const float* xq_train) override;
+
+    void reset() override;
+
+    void serialize_training_data(std::ostream& out) const;
+    void deserialize_training_data(std::istream& in);
+
+    size_t leanvec_d = 0;
+
+    /* Training information */
+    svs_runtime::LeanVecTrainingData* training_data{nullptr};
+
+   protected:
+    void create_impl(idx_t n, const float* x) override;
+};
+
+} // namespace faiss

data/vendor/faiss/faiss/svs/IndexSVSVamana.cpp

@@ -31,6 +31,7 @@
 #include <svs/runtime/vamana_index.h>
 
 #include <cstddef>
+#include <cstring>
 #include <numeric>
 #include <span>
 #include <type_traits>
@@ -114,13 +115,32 @@ void IndexSVSVamana::add(idx_t n, const float* x) {
     if (!status.ok()) {
         FAISS_THROW_MSG(status.message());
     }
+
+    if (stored_vectors_valid) {
+        size_t prev = static_cast<size_t>(ntotal) * d;
+        stored_vectors.resize(prev + static_cast<size_t>(n) * d);
+        std::memcpy(stored_vectors.data() + prev, x, sizeof(float) * n * d);
+    }
     ntotal += n;
 }
 
+void IndexSVSVamana::reconstruct(idx_t key, float* recons) const {
+    FAISS_THROW_IF_NOT_MSG(
+            key >= 0 && key < ntotal,
+            "IndexSVSVamana::reconstruct: key out of range");
+    FAISS_THROW_IF_NOT_MSG(
+            stored_vectors_valid && !stored_vectors.empty(),
+            "IndexSVSVamana::reconstruct: stored_vectors unavailable "
+            "(invalidated by remove_ids or not restored after deserialization)");
+    std::memcpy(recons, stored_vectors.data() + key * d, sizeof(float) * d);
+}
+
 void IndexSVSVamana::reset() {
     if (impl) {
         impl->reset();
     }
+    stored_vectors.clear();
+    stored_vectors_valid = true;
     is_trained = false;
     ntotal = 0;
 }
@@ -189,6 +209,8 @@ size_t IndexSVSVamana::remove_ids(const IDSelector& sel) {
     size_t removed = 0;
     auto Status = impl->remove_selected(&removed, id_filter);
     ntotal -= removed;
+    stored_vectors.clear();
+    stored_vectors_valid = false;
     return removed;
 }
 
@@ -236,10 +258,12 @@ void IndexSVSVamana::deserialize_impl(std::istream& in) {
     FAISS_THROW_IF_MSG(impl, "Cannot deserialize: SVS index already loaded.");
     auto svs_metric = to_svs_metric(metric_type);
     auto svs_storage_kind = to_svs_storage_kind(storage_kind);
-    auto status = impl->load(&impl, in, svs_metric, svs_storage_kind);
+    auto status = svs_runtime::DynamicVamanaIndex::load(
+            &impl, in, svs_metric, svs_storage_kind);
     if (!status.ok()) {
         FAISS_THROW_MSG(status.message());
     }
+    FAISS_THROW_IF_NOT_MSG(impl, "Failed to load SVS Vamana index.");
 }
 
 } // namespace faiss

data/vendor/faiss/faiss/svs/IndexSVSVamana.h

@@ -27,8 +27,11 @@
 #include <faiss/svs/IndexSVSFaissUtils.h>
 
 #include <svs/runtime/api_defs.h>
+#include <svs/runtime/dynamic_vamana_index.h>
 
 #include <iostream>
+#include <type_traits>
+#include <vector>
 
 namespace faiss {
 
@@ -45,9 +48,11 @@ enum SVSStorageKind {
     SVS_LVQ4x0,
     SVS_LVQ4x4,
     SVS_LVQ4x8,
+    SVS_LVQ8x0,
     SVS_LeanVec4x4,
     SVS_LeanVec4x8,
     SVS_LeanVec8x8,
+    SVS_count,
 };
 
 inline svs_runtime::StorageKind to_svs_storage_kind(SVSStorageKind kind) {
@@ -64,6 +69,8 @@ inline svs_runtime::StorageKind to_svs_storage_kind(SVSStorageKind kind) {
             return svs_runtime::StorageKind::LVQ4x4;
         case SVS_LVQ4x8:
             return svs_runtime::StorageKind::LVQ4x8;
+        case SVS_LVQ8x0:
+            return svs_runtime::StorageKind::LVQ8x0;
         case SVS_LeanVec4x4:
             return svs_runtime::StorageKind::LeanVec4x4;
         case SVS_LeanVec4x8:
@@ -71,7 +78,9 @@ inline svs_runtime::StorageKind to_svs_storage_kind(SVSStorageKind kind) {
         case SVS_LeanVec8x8:
             return svs_runtime::StorageKind::LeanVec8x8;
         default:
-            FAISS_ASSERT(!"not supported SVS storage kind");
+            FAISS_THROW_FMT(
+                    "SVSStorageKind (%d) not supported",
+                    static_cast<std::underlying_type_t<SVSStorageKind>>(kind));
     }
 }
 
@@ -85,7 +94,7 @@ struct IndexSVSVamana : Index {
     size_t max_candidate_pool_size = 200;
     bool use_full_search_history = true;
 
-    SVSStorageKind storage_kind;
+    SVSStorageKind storage_kind = SVS_FP32;
 
     IndexSVSVamana();
 
@@ -103,6 +112,8 @@ struct IndexSVSVamana : Index {
 
     void add(idx_t n, const float* x) override;
 
+    void reconstruct(idx_t key, float* recons) const override;
+
     void search(
             idx_t n,
             const float* x,
@@ -129,6 +140,12 @@ struct IndexSVSVamana : Index {
     /* The actual SVS implementation */
     svs_runtime::DynamicVamanaIndex* impl{nullptr};
 
+    // The SVS runtime API does not expose vector retrieval, so we keep a copy
+    // of added vectors to support reconstruct(). When used as a coarse
+    // quantizer this holds only nlist centroids.
+    std::vector<float> stored_vectors;
+    bool stored_vectors_valid{true};
+
    protected:
     /* Initializes the implementation*/
     virtual void create_impl();

data/vendor/faiss/faiss/svs/IndexSVSVamanaLVQ.h

@@ -34,7 +34,7 @@ struct IndexSVSVamanaLVQ : IndexSVSVamana {
             idx_t d,
             size_t degree,
             MetricType metric = METRIC_L2,
-            SVSStorageKind storage = SVSStorageKind::SVS_LVQ4x4);
+            SVSStorageKind storage = SVSStorageKind::SVS_LVQ4x0);
 
     ~IndexSVSVamanaLVQ() override = default;
 };

data/vendor/faiss/faiss/svs/IndexSVSVamanaLeanVec.cpp

@@ -56,7 +56,7 @@ IndexSVSVamanaLeanVec::~IndexSVSVamanaLeanVec() {
         FAISS_ASSERT(status.ok());
         training_data = nullptr;
     }
-    IndexSVSVamana::~IndexSVSVamana();
+    // Base class destructor handles impl cleanup
 }
 
 void IndexSVSVamanaLeanVec::add(idx_t n, const float* x) {
@@ -66,6 +66,14 @@ void IndexSVSVamanaLeanVec::add(idx_t n, const float* x) {
 }
 
 void IndexSVSVamanaLeanVec::train(idx_t n, const float* x) {
+    train_with_queries(n, x, 0, nullptr);
+}
+
+void IndexSVSVamanaLeanVec::train_with_queries(
+        idx_t n,
+        const float* x,
+        idx_t n_train_q,
+        const float* queries) {
     FAISS_THROW_IF_MSG(
             training_data || impl, "Index already trained or contains data.");
 
@@ -74,7 +82,7 @@ void IndexSVSVamanaLeanVec::train(idx_t n, const float* x) {
             "LVQ/LeanVec support not available on this platform or build");
 
     auto status = svs_runtime::LeanVecTrainingData::build(
-            &training_data, d, n, x, leanvec_d);
+            &training_data, d, n, x, n_train_q, queries, leanvec_d);
     if (!status.ok()) {
         FAISS_THROW_MSG(status.message());
     }
@@ -83,6 +91,15 @@ void IndexSVSVamanaLeanVec::train(idx_t n, const float* x) {
     is_trained = true;
 }
 
+void IndexSVSVamanaLeanVec::reset() {
+    if (training_data) {
+        auto status = svs_runtime::LeanVecTrainingData::destroy(training_data);
+        FAISS_ASSERT(status.ok());
+        training_data = nullptr;
+    }
+    IndexSVSVamana::reset();
+}
+
 void IndexSVSVamanaLeanVec::serialize_training_data(std::ostream& out) const {
     FAISS_THROW_IF_NOT_MSG(
             training_data, "Cannot serialize: Training data not initialized.");

data/vendor/faiss/faiss/svs/IndexSVSVamanaLeanVec.h

@@ -41,8 +41,22 @@ struct IndexSVSVamanaLeanVec : IndexSVSVamana {
 
     void add(idx_t n, const float* x) override;
 
+    /* Default train assumes in-distribution data */
     void train(idx_t n, const float* x) override;
 
+    /* Generic train with out-of-distribution parameters.
+     * Out-of-distribution (OOD) means database vectors and queries _can_ be
+     * sampled from different distributions (e.g., cross-modal). More details in
+     * the original publication, arXiv:2312.16335.
+     */
+    void train_with_queries(
+            idx_t n,
+            const float* x,
+            idx_t n_train_q,
+            const float* xq_train) override;
+
+    void reset() override;
+
     void serialize_training_data(std::ostream& out) const;
     void deserialize_training_data(std::istream& in);
 

data/vendor/faiss/faiss/utils/Heap.cpp

@@ -17,7 +17,7 @@ namespace faiss {
 template <typename C>
 void HeapArray<C>::heapify() {
 #pragma omp parallel for
-    for (int64_t j = 0; j < nh; j++) {
+    for (int64_t j = 0; j < static_cast<int64_t>(nh); j++) {
         heap_heapify<C>(k, val + j * k, ids + j * k);
     }
 }
@@ -25,7 +25,7 @@ void HeapArray<C>::heapify() {
 template <typename C>
 void HeapArray<C>::reorder() {
 #pragma omp parallel for
-    for (int64_t j = 0; j < nh; j++) {
+    for (int64_t j = 0; j < static_cast<int64_t>(nh); j++) {
         heap_reorder<C>(k, val + j * k, ids + j * k);
     }
 }
@@ -37,7 +37,7 @@ void HeapArray<C>::addn(size_t nj, const T* vin, TI j0, size_t i0, int64_t ni) {
     }
     assert(i0 >= 0 && i0 + ni <= nh);
 #pragma omp parallel for if (ni * nj > 100000)
-    for (int64_t i = i0; i < i0 + ni; i++) {
+    for (int64_t i = i0; i < static_cast<int64_t>(i0) + ni; i++) {
         T* __restrict simi = get_val(i);
         TI* __restrict idxi = get_ids(i);
         const T* ip_line = vin + (i - i0) * nj;
@@ -68,7 +68,7 @@ void HeapArray<C>::addn_with_ids(
     }
     assert(i0 >= 0 && i0 + ni <= nh);
 #pragma omp parallel for if (ni * nj > 100000)
-    for (int64_t i = i0; i < i0 + ni; i++) {
+    for (int64_t i = i0; i < static_cast<int64_t>(i0) + ni; i++) {
         T* __restrict simi = get_val(i);
         TI* __restrict idxi = get_ids(i);
         const T* ip_line = vin + (i - i0) * nj;
@@ -96,7 +96,7 @@ void HeapArray<C>::addn_query_subset_with_ids(
         id_stride = nj;
     }
 #pragma omp parallel for if (nsubset * nj > 100000)
-    for (int64_t si = 0; si < nsubset; si++) {
+    for (int64_t si = 0; si < static_cast<int64_t>(nsubset); si++) {
         TI i = subset[si];
         T* __restrict simi = get_val(i);
         TI* __restrict idxi = get_ids(i);
@@ -115,7 +115,7 @@ void HeapArray<C>::addn_query_subset_with_ids(
 template <typename C>
 void HeapArray<C>::per_line_extrema(T* out_val, TI* out_ids) const {
 #pragma omp parallel for if (nh * k > 100000)
-    for (int64_t j = 0; j < nh; j++) {
+    for (int64_t j = 0; j < static_cast<int64_t>(nh); j++) {
         int64_t imin = -1;
         typename C::T xval = C::Crev::neutral();
         const typename C::T* x_ = val + j * k;
@@ -187,7 +187,7 @@ void merge_knn_results(
         std::vector<distance_t> buf2(nshard);
         distance_t* heap_vals = buf2.data();
 #pragma omp for
-        for (long i = 0; i < n; i++) {
+        for (long i = 0; i < static_cast<long>(n); i++) {
             // the heap maps values to the shard where they are
             // produced.
             const distance_t* D_in = all_distances + i * k;
@@ -195,7 +195,7 @@ void merge_knn_results(
             int heap_size = 0;
 
             // push the first element of each shard (if not -1)
-            for (long s = 0; s < nshard; s++) {
+            for (long s = 0; s < static_cast<long>(nshard); s++) {
                 pointer[s] = 0;
                 if (I_in[stride * s] >= 0) {
                     heap_push<C>(
@@ -210,7 +210,7 @@ void merge_knn_results(
             distance_t* D = distances + i * k;
             idx_t* I = labels + i * k;
 
-            int j;
+            size_t j;
             for (j = 0; j < k && heap_size > 0; j++) {
                 // pop element from best shard
                 int s = shard_ids[0]; // top of heap
@@ -221,7 +221,7 @@ void merge_knn_results(
                 // pop from shard, advance pointer for this shard
                 heap_pop<C>(heap_size--, heap_vals, shard_ids);
                 p++;
-                if (p < k && I_in[stride * s + p] >= 0) {
+                if (static_cast<size_t>(p) < k && I_in[stride * s + p] >= 0) {
                     heap_push<C>(
                             ++heap_size,
                             heap_vals,
@@ -254,4 +254,50 @@ INSTANTIATE(CMax, float);
 INSTANTIATE(CMin, int32_t);
 INSTANTIATE(CMax, int32_t);
 
+/**********************************************************
+ * reorder_2_heaps
+ **********************************************************/
+
+template <class C>
+void reorder_2_heaps(
+        int64_t n,
+        int64_t k,
+        typename C::TI* __restrict labels,
+        float* __restrict distances,
+        int64_t k_base,
+        const typename C::TI* __restrict base_labels,
+        const float* __restrict base_distances) {
+#pragma omp parallel for if (n > 1)
+    for (int64_t i = 0; i < n; i++) {
+        typename C::TI* idxo = labels + i * k;
+        float* diso = distances + i * k;
+        const typename C::TI* idxi = base_labels + i * k_base;
+        const float* disi = base_distances + i * k_base;
+
+        heap_heapify<C>(k, diso, idxo, disi, idxi, k);
+        if (k_base != k) { // add remaining elements
+            heap_addn<C>(k, diso, idxo, disi + k, idxi + k, k_base - k);
+        }
+        heap_reorder<C>(k, diso, idxo);
+    }
+}
+
+template void reorder_2_heaps<CMax<float, int64_t>>(
+        int64_t n,
+        int64_t k,
+        int64_t* __restrict labels,
+        float* __restrict distances,
+        int64_t k_base,
+        const int64_t* __restrict base_labels,
+        const float* __restrict base_distances);
+
+template void reorder_2_heaps<CMin<float, int64_t>>(
+        int64_t n,
+        int64_t k,
+        int64_t* __restrict labels,
+        float* __restrict distances,
+        int64_t k_base,
+        const int64_t* __restrict base_labels,
+        const float* __restrict base_distances);
+
 } // namespace faiss

data/vendor/faiss/faiss/utils/Heap.h

@@ -642,6 +642,27 @@ void merge_knn_results(
         typename C::T* distances,
         idx_t* labels);
 
+/** Reduces k_base pairs (base_labels, base_distances) into k pairs
+ * (labels, distances). The function is used for the refining process.
+ *
+ * @param n              number of vectors to process
+ * @param k              number of output nearest neighbors per vector
+ * @param labels         output labels, size (n, k)
+ * @param distances      output distances, size (n, k)
+ * @param k_base         number of input nearest neighbors per vector
+ * @param base_labels    input labels, size (n, k_base)
+ * @param base_distances input distances, size (n, k_base)
+ */
+template <class C>
+void reorder_2_heaps(
+        int64_t n,
+        int64_t k,
+        typename C::TI* __restrict labels,
+        float* __restrict distances,
+        int64_t k_base,
+        const typename C::TI* __restrict base_labels,
+        const float* __restrict base_distances);
+
 } // namespace faiss
 
 #endif /* FAISS_Heap_h */