faiss 0.6.4 → 0.6.5

data/vendor/faiss/faiss/impl/DistanceComputer.h

@@ -41,17 +41,7 @@ struct DistanceComputer {
             float& dis0,
             float& dis1,
             float& dis2,
-            float& dis3) {
-        // compute first, assign next
-        const float d0 = this->operator()(idx0);
-        const float d1 = this->operator()(idx1);
-        const float d2 = this->operator()(idx2);
-        const float d3 = this->operator()(idx3);
-        dis0 = d0;
-        dis1 = d1;
-        dis2 = d2;
-        dis3 = d3;
-    }
+            float& dis3);
 
     /// compute distance between two stored vectors
     virtual float symmetric_dis(idx_t i, idx_t j) = 0;
@@ -181,9 +171,7 @@ struct FlatCodesDistanceComputer : DistanceComputer {
     virtual float partial_dot_product(
             const idx_t /* i */,
             const uint32_t /* offset */,
-            const uint32_t /* num_components */) {
-        FAISS_THROW_MSG("partial_dot_product not implemented");
-    }
+            const uint32_t /* num_components */);
 
     /// compute distance of current query to an encoded vector
     virtual float distance_to_code(const uint8_t* code) = 0;
@@ -214,21 +202,7 @@ struct FlatCodesDistanceComputer : DistanceComputer {
             float& dp2,
             float& dp3,
             const uint32_t offset,
-            const uint32_t num_components) {
-        // default implementation for correctness
-        const float d0 =
-                this->partial_dot_product(idx0, offset, num_components);
-        const float d1 =
-                this->partial_dot_product(idx1, offset, num_components);
-        const float d2 =
-                this->partial_dot_product(idx2, offset, num_components);
-        const float d3 =
-                this->partial_dot_product(idx3, offset, num_components);
-        dp0 = d0;
-        dp1 = d1;
-        dp2 = d2;
-        dp3 = d3;
-    }
+            const uint32_t num_components);
 
     virtual ~FlatCodesDistanceComputer() override {}
 };

data/vendor/faiss/faiss/impl/HNSW.cpp

@@ -1568,6 +1568,8 @@ void reservePriorityQueue(
 /// Templated body of `search_from_candidate_unbounded`. The choice of
 /// max-heap vs min-heap for both `top_candidates` and `candidates` is
 /// derived from C via `TopCandidatesQueue` / `CandidatesQueue`.
+/// `sel` filters the results only. The traversal still visits the rejected
+/// nodes.
 template <typename VTType, class C>
 TopCandidatesQueue<C> search_from_candidate_unbounded_fixVT(
         const HNSW& hnsw,
@@ -1575,16 +1577,48 @@ TopCandidatesQueue<C> search_from_candidate_unbounded_fixVT(
         DistanceComputer& qdis,
         int ef,
         VTType& vt,
-        HNSWStats& stats) {
+        HNSWStats& stats,
+        const IDSelector* sel) {
+    const auto max_size = static_cast<size_t>(ef);
+
     int ndis = 0;
+    // `top_candidates` bounds the traversal and keeps the nodes that `sel`
+    // rejects, because a rejected node is still a valid hop towards an
+    // accepted one.
     TopCandidatesQueue<C> top_candidates;
     reservePriorityQueue(top_candidates, ef);
 
+    // `selected_candidates` keeps only the accepted nodes, and
+    // the function returns it instead when `sel` is set.
+    TopCandidatesQueue<C> selected_candidates;
+    if (sel) {
+        reservePriorityQueue(selected_candidates, ef);
+    }
+
     CandidatesQueue<C> candidates;
     reservePriorityQueue(candidates, ef);
 
+    // Adds a node to a queue of the best `max_size` nodes. The size test
+    // comes first because `top()` needs a queue that is not empty.
+    auto push_bounded = [max_size](
+                                TopCandidatesQueue<C>& queue,
+                                const float dis,
+                                const size_t idx) {
+        if (queue.size() >= max_size && !C::cmp(queue.top().first, dis)) {
+            return false;
+        }
+        queue.emplace(dis, idx);
+        if (queue.size() > max_size) {
+            queue.pop();
+        }
+        return true;
+    };
+
     top_candidates.push(node);
     candidates.push(node);
+    if (sel && sel->is_member(node.second)) {
+        selected_candidates.push(node);
+    }
 
     vt.set(node.second);
 
@@ -1619,14 +1653,11 @@ TopCandidatesQueue<C> search_from_candidate_unbounded_fixVT(
         size_t saved_j[4];
 
         auto add_to_heap = [&](const size_t idx, const float dis) {
-            if (C::cmp(top_candidates.top().first, dis) ||
-                top_candidates.size() < static_cast<size_t>(ef)) {
+            if (sel && sel->is_member(idx)) {
+                push_bounded(selected_candidates, dis, idx);
+            }
+            if (push_bounded(top_candidates, dis, idx)) {
                 candidates.emplace(dis, idx);
-                top_candidates.emplace(dis, idx);
-
-                if (top_candidates.size() > static_cast<size_t>(ef)) {
-                    top_candidates.pop();
-                }
             }
         };
 
@@ -1674,6 +1705,9 @@ TopCandidatesQueue<C> search_from_candidate_unbounded_fixVT(
     }
     stats.ndis += ndis;
 
+    if (sel) {
+        return selected_candidates;
+    }
     return top_candidates;
 }
 
@@ -1693,7 +1727,7 @@ std::priority_queue<HNSW::Node> hnsw_detail::search_from_candidate_unbounded(
     using C = HNSW::C_distance;
     auto call = [&]<typename VTType>(VTType& vt_concrete) {
         return search_from_candidate_unbounded_fixVT<VTType, C>(
-                hnsw, node, qdis, ef, vt_concrete, stats);
+                hnsw, node, qdis, ef, vt_concrete, stats, nullptr);
     };
     if (VisitedTableVector* vtv = dynamic_cast<VisitedTableVector*>(vt)) {
         return call(*vtv);
@@ -1792,6 +1826,7 @@ HNSWStats search_impl(
                 hnsw.search_method == HNSW::SM_DEFAULT ||
                         hnsw.search_method == HNSW::SM_PANORAMA,
                 "invalid HNSW search method");
+        const IDSelector* sel = params ? params->sel : nullptr;
         auto call = [&]<typename VTType>(VTType& vt_concrete) {
             return search_from_candidate_unbounded_fixVT<VTType, C>(
                     hnsw,
@@ -1799,7 +1834,8 @@ HNSWStats search_impl(
                     qdis,
                     ef,
                     vt_concrete,
-                    stats);
+                    stats,
+                    sel);
         };
         TopCandidatesQueue<C> top_candidates;
         if (VisitedTableVector* vtv = dynamic_cast<VisitedTableVector*>(&vt)) {

data/vendor/faiss/faiss/impl/HNSW.h

@@ -296,6 +296,76 @@ void hnsw_add_vertices_deterministic(
         const std::function<void(DistanceComputer&, HNSW::storage_idx_t)>&
                 set_query);
 
+/// Parameters of add_level_0_links_from_sample() of IndexHNSW and
+/// IndexBinaryHNSW.
+struct HNSWLevel0LinkParameters {
+    /// The search that the new links are for. nullptr: the search settings
+    /// of the index, which the IVF assignment with a coarse quantizer uses.
+    const SearchParametersHNSW* search_params = nullptr;
+    /// A sink receives at least this factor times the mean number of
+    /// vectors per node.
+    float min_mean_factor = 3.0f;
+    /// For at least this fraction of the vectors that the search sends to a
+    /// sink, an exact search finds a nearer node.
+    float min_closer_fraction = 0.5f;
+    /// Maximum number of level-0 links added to one sink in one round.
+    int max_links_per_sink = 32;
+    /// Maximum number of rounds that find sinks and add links.
+    int max_rounds = 3;
+};
+
+struct HNSWLevel0LinkStats {
+    /// Sinks found in the first round.
+    size_t n_sinks_found = 0;
+    /// Sinks found after the last round.
+    size_t n_sinks_remaining = 0;
+    /// Level-0 links written into free slots.
+    size_t n_links_added = 0;
+    /// Level-0 links that replaced a link.
+    size_t n_links_replaced = 0;
+    /// Rounds that changed links.
+    int n_rounds = 0;
+    /// Exact searches with the storage, in all rounds.
+    size_t n_exact_searches = 0;
+};
+
+/** The two searches of the sample vectors that
+ * hnsw_add_level_0_links_from_sample() compares. IndexHNSW and IndexBinaryHNSW
+ * implement them for their vector types. */
+struct HNSWLinkSampleSearch {
+    /// number of sample vectors
+    idx_t n = 0;
+
+    virtual ~HNSWLinkSampleSearch() = default;
+
+    /// HNSW search (k = 1) of the m sample vectors `ids` (null: the first m).
+    virtual void search(
+            const SearchParametersHNSW* params,
+            size_t m,
+            const idx_t* ids,
+            idx_t* labels) const = 0;
+
+    /** Exact search (k = 1) of the m sample vectors `ids`. nearer[j] is 1 when
+     * exact[j] is nearer than found[j] by more than rounding differences. */
+    virtual void search_exact(
+            size_t m,
+            const idx_t* ids,
+            const idx_t* found,
+            idx_t* exact,
+            uint8_t* nearer) const = 0;
+};
+
+/** add_level_0_links_from_sample() of IndexHNSW and IndexBinaryHNSW.
+ * `dis` computes distances between nodes. `similarity` is true when a larger
+ * distance is nearer. */
+HNSWLevel0LinkStats hnsw_add_level_0_links_from_sample(
+        HNSW& hnsw,
+        size_t ntotal,
+        DistanceComputer& dis,
+        bool similarity,
+        const HNSWLinkSampleSearch& sample,
+        const HNSWLevel0LinkParameters& params);
+
 struct HNSWStats {
     size_t n1 = 0; /// number of vectors searched
     size_t n2 =
@@ -315,6 +385,8 @@ struct HNSWStats {
         ndis += other.ndis;
         nhops += other.nhops;
     }
+
+    void combine_atomic(const HNSWStats& other);
 };
 
 // global var that collects them all

data/vendor/faiss/faiss/impl/LocalSearchQuantizer.cpp

@@ -23,12 +23,9 @@
 #include <faiss/utils/utils.h>
 
 #include <faiss/impl/approx_topk/approx_topk.h>
+#include <faiss/impl/platform_macros.h>
 
-// this is needed for prefetching
-
-#ifdef COMPILE_SIMD_AVX2
-#include <xmmintrin.h>
-#endif
+#include <faiss/utils/prefetch.h>
 
 extern "C" {
 // LU decomposition of a general matrix
@@ -602,7 +599,7 @@ void LocalSearchQuantizer::icm_encode_step(
 
     // Resolve SIMD level once, not per iteration of the n × n_iters × M loop.
     with_simd_level_256bit([&]<SIMDLevel SL>() {
-#pragma omp parallel for schedule(dynamic)
+#pragma omp parallel for schedule(FAISS_OMP_DYNAMIC)
         for (int64_t i = 0; i < static_cast<int64_t>(n); i++) {
             std::vector<float> objs(K);
 
@@ -622,23 +619,13 @@ void LocalSearchQuantizer::icm_encode_step(
                             continue;
                         }
 
-#ifdef COMPILE_SIMD_AVX2
-                        // TODO: add platform-independent compiler-independent
-                        // prefetch utilities.
                         if (other_m + 1 < M) {
                             // do a single prefetch
                             int32_t code2 = codes[i * M + other_m + 1];
-                            // for (int32_t code = 0; code < K; code += 64) {
-                            int32_t code = 0;
-                            {
-                                size_t binary_idx = (other_m + 1) * M * K * K +
-                                        m * K * K + code2 * K + code;
-                                _mm_prefetch(
-                                        (const char*)(binaries + binary_idx),
-                                        _MM_HINT_T0);
-                            }
+                            size_t binary_idx = (other_m + 1) * M * K * K +
+                                    m * K * K + code2 * K;
+                            prefetch_L1(binaries + binary_idx);
                         }
-#endif
 
                         for (size_t code = 0; code < K; code++) {
                             int32_t code2 = codes[i * M + other_m];

data/vendor/faiss/faiss/impl/NNDescent.cpp

@@ -14,6 +14,7 @@
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/DistanceComputer.h>
 #include <faiss/impl/VisitedTable.h>
+#include <faiss/impl/platform_macros.h>
 
 namespace faiss {
 
@@ -200,7 +201,7 @@ void NNDescent::join(DistanceComputer& qdis) {
     idx_t check_period = InterruptCallback::get_period_hint(d * search_L);
     for (idx_t i0 = 0; i0 < (idx_t)ntotal; i0 += check_period) {
         idx_t i1 = std::min(i0 + check_period, (idx_t)ntotal);
-#pragma omp parallel for default(shared) schedule(dynamic, 100)
+#pragma omp parallel for default(shared) schedule(FAISS_OMP_DYNAMIC, 100)
         for (idx_t n = i0; n < i1; n++) {
             graph[n].join([&](int i, int j) {
                 if (i != j) {

data/vendor/faiss/faiss/impl/NSG.cpp

@@ -15,6 +15,7 @@
 
 #include <faiss/impl/DistanceComputer.h>
 #include <faiss/impl/VisitedTable.h>
+#include <faiss/impl/platform_macros.h>
 
 namespace faiss {
 
@@ -352,7 +353,7 @@ void NSG::link(
         std::unique_ptr<DistanceComputer> dis(
                 storage_distance_computer(storage));
 
-#pragma omp for schedule(dynamic, 100)
+#pragma omp for schedule(FAISS_OMP_DYNAMIC, 100)
         for (int i = 0; i < ntotal; i++) {
             storage->reconstruct(i, vec.get());
             dis->set_query(vec.get());
@@ -375,7 +376,7 @@ void NSG::link(
         std::unique_ptr<DistanceComputer> dis(
                 storage_distance_computer(storage));
 
-#pragma omp for schedule(dynamic, 100)
+#pragma omp for schedule(FAISS_OMP_DYNAMIC, 100)
         for (int i = 0; i < ntotal; ++i) {
             add_reverse_links(i, locks, *dis, graph);
         }

data/vendor/faiss/faiss/impl/Panorama.h

@@ -14,6 +14,7 @@
 #include <faiss/impl/IDSelector.h>
 #include <faiss/impl/PanoramaStats.h>
 #include <faiss/utils/distances.h>
+#include <faiss/utils/simd_levels.h>
 
 #include <algorithm>
 #include <cstddef>
@@ -21,7 +22,7 @@
 #include <cstring>
 #include <vector>
 
-#if defined(COMPILE_SIMD_AVX2) && defined(__AVX2__) && defined(__BMI2__)
+#if FAISS_SIMD_TU(AVX2) && defined(__BMI2__)
 #include <immintrin.h>
 #endif
 
@@ -154,7 +155,7 @@ inline size_t compact_active_kernel(
     size_t next_active = 0;
     size_t i = 0;
 
-#if defined(COMPILE_SIMD_AVX2) && defined(__AVX2__) && defined(__BMI2__)
+#if FAISS_SIMD_TU(AVX2) && defined(__BMI2__)
     for (; i + 8 <= num_active; i += 8) {
         uint64_t bytes;
         memcpy(&bytes, &active_byteset[i], 8);