faiss 0.6.0 → 0.6.1

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

@@ -17,7 +17,6 @@
 #include <faiss/impl/DistanceComputer.h>
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/maybe_owned_vector.h>
-#include <faiss/impl/platform_macros.h>
 #include <faiss/utils/Heap.h>
 #include <faiss/utils/random.h>
 
@@ -26,6 +25,8 @@ namespace faiss {
 // Forward declarations to avoid circular dependency.
 struct IndexHNSW;
 struct IndexHNSWFlatPanorama;
+struct MinimaxHeap;
+class LockVector;
 
 /** Implementation of the Hierarchical Navigable Small World
  * datastructure.
@@ -64,37 +65,11 @@ struct HNSW {
 
     typedef std::pair<float, storage_idx_t> Node;
 
-    /** Heap structure that allows fast access and updates.
-     */
-    struct MinimaxHeap {
-        int n;
-        int k;
-        int nvalid;
-
-        std::vector<storage_idx_t> ids;
-        std::vector<float> dis;
-        typedef faiss::CMax<float, storage_idx_t> HC;
-
-        explicit MinimaxHeap(int n) : n(n), k(0), nvalid(0), ids(n), dis(n) {}
-
-        void push(storage_idx_t i, float v);
-
-        float max() const;
-
-        int size() const;
-
-        void clear();
-
-        int pop_min(float* vmin_out = nullptr);
-
-        int count_below(float thresh);
-    };
-
     /// to sort pairs of (id, distance) from nearest to farthest or the reverse
     struct NodeDistCloser {
         float d;
         int id;
-        NodeDistCloser(float d, int id) : d(d), id(id) {}
+        NodeDistCloser(float d_in, int id_in) : d(d_in), id(id_in) {}
         bool operator<(const NodeDistCloser& obj1) const {
             return d < obj1.d;
         }
@@ -103,7 +78,7 @@ struct HNSW {
     struct NodeDistFarther {
         float d;
         int id;
-        NodeDistFarther(float d, int id) : d(d), id(id) {}
+        NodeDistFarther(float d_in, int id_in) : d(d_in), id(id_in) {}
         bool operator<(const NodeDistFarther& obj1) const {
             return d > obj1.d;
         }
@@ -142,6 +117,10 @@ struct HNSW {
     /// expansion factor at search time
     int efSearch = 16;
 
+    /// when pruning, leave room for more neighbors to avoid O(n^2)
+    /// costs and lock contention on frequently-pruned nodes.
+    float prune_headroom = 0.2f;
+
     /// during search: do we check whether the next best distance is good
     /// enough?
     bool check_relative_distance = true;
@@ -191,7 +170,7 @@ struct HNSW {
             storage_idx_t nearest,
             float d_nearest,
             int level,
-            omp_lock_t* locks,
+            LockVector& locks,
             VisitedTable& vt,
             bool keep_max_size_level0 = false);
 
@@ -201,7 +180,7 @@ struct HNSW {
             DistanceComputer& ptdis,
             int pt_level,
             int pt_id,
-            std::vector<omp_lock_t>& locks,
+            LockVector& locks,
             VisitedTable& vt,
             bool keep_max_size_level0 = false);
 
@@ -241,7 +220,7 @@ struct HNSW {
             DistanceComputer& qdis,
             std::priority_queue<NodeDistFarther>& input,
             std::vector<NodeDistFarther>& output,
-            int max_size,
+            size_t max_size,
             bool keep_max_size_level0 = false);
 
     void permute_entries(const idx_t* map);
@@ -275,7 +254,7 @@ int search_from_candidates(
         const HNSW& hnsw,
         DistanceComputer& qdis,
         ResultHandler& res,
-        HNSW::MinimaxHeap& candidates,
+        MinimaxHeap& candidates,
         VisitedTable& vt,
         HNSWStats& stats,
         int level,
@@ -291,7 +270,7 @@ int search_from_candidates_panorama(
         const IndexHNSW* index,
         DistanceComputer& qdis,
         ResultHandler& res,
-        HNSW::MinimaxHeap& candidates,
+        MinimaxHeap& candidates,
         VisitedTable& vt,
         HNSWStats& stats,
         int level,

data/vendor/faiss/faiss/impl/IDSelector.cpp

@@ -14,8 +14,11 @@ namespace faiss {
  * IDSelectorRange
  ***********************************************************************/
 
-IDSelectorRange::IDSelectorRange(idx_t imin, idx_t imax, bool assume_sorted)
-        : imin(imin), imax(imax), assume_sorted(assume_sorted) {}
+IDSelectorRange::IDSelectorRange(
+        idx_t imin_in,
+        idx_t imax_in,
+        bool assume_sorted_in)
+        : imin(imin_in), imax(imax_in), assume_sorted(assume_sorted_in) {}
 
 bool IDSelectorRange::is_member(idx_t id) const {
     return id >= imin && id < imax;
@@ -67,10 +70,11 @@ void IDSelectorRange::find_sorted_ids_bounds(
  * IDSelectorArray
  ***********************************************************************/
 
-IDSelectorArray::IDSelectorArray(size_t n, const idx_t* ids) : n(n), ids(ids) {}
+IDSelectorArray::IDSelectorArray(size_t n_in, const idx_t* ids_in)
+        : n(n_in), ids(ids_in) {}
 
 bool IDSelectorArray::is_member(idx_t id) const {
-    for (idx_t i = 0; i < n; i++) {
+    for (size_t i = 0; i < n; i++) {
         if (ids[i] == id) {
             return true;
         }
@@ -84,15 +88,15 @@ bool IDSelectorArray::is_member(idx_t id) const {
 
 IDSelectorBatch::IDSelectorBatch(size_t n, const idx_t* indices) {
     nbits = 0;
-    while (n > ((idx_t)1 << nbits)) {
+    while (n > (size_t{1} << nbits)) {
         nbits++;
     }
     nbits += 5;
     // for n = 1M, nbits = 25 is optimal, see P56659518
 
     mask = ((idx_t)1 << nbits) - 1;
-    bloom.resize((idx_t)1 << (nbits - 3), 0);
-    for (idx_t i = 0; i < n; i++) {
+    bloom.resize(size_t{1} << (nbits - 3), 0);
+    for (size_t i = 0; i < n; i++) {
         idx_t id = indices[i];
         set.insert(id);
         id &= mask;
@@ -101,9 +105,9 @@ IDSelectorBatch::IDSelectorBatch(size_t n, const idx_t* indices) {
 }
 
 bool IDSelectorBatch::is_member(idx_t i) const {
-    long im = i & mask;
+    idx_t im = i & mask;
     if (!(bloom[im >> 3] & (1 << (im & 7)))) {
-        return 0;
+        return false;
     }
     return set.count(i);
 }
@@ -112,8 +116,8 @@ bool IDSelectorBatch::is_member(idx_t i) const {
  * IDSelectorBitmap
  ***********************************************************************/
 
-IDSelectorBitmap::IDSelectorBitmap(size_t n, const uint8_t* bitmap)
-        : n(n), bitmap(bitmap) {}
+IDSelectorBitmap::IDSelectorBitmap(size_t n_in, const uint8_t* bitmap_in)
+        : n(n_in), bitmap(bitmap_in) {}
 
 bool IDSelectorBitmap::is_member(idx_t ii) const {
     uint64_t i = ii;

data/vendor/faiss/faiss/impl/IDSelector.h

@@ -116,7 +116,7 @@ struct IDSelectorBitmap : IDSelector {
 /** reverts the membership test of another selector */
 struct IDSelectorNot : IDSelector {
     const IDSelector* sel;
-    explicit IDSelectorNot(const IDSelector* sel) : sel(sel) {}
+    explicit IDSelectorNot(const IDSelector* sel_) : sel(sel_) {}
     bool is_member(idx_t id) const final {
         return !sel->is_member(id);
     }
@@ -125,7 +125,7 @@ struct IDSelectorNot : IDSelector {
 
 /// selects all entries (useful for benchmarking)
 struct IDSelectorAll : IDSelector {
-    bool is_member(idx_t id) const final {
+    bool is_member(idx_t /* id */) const final {
         return true;
     }
     virtual ~IDSelectorAll() {}
@@ -136,8 +136,8 @@ struct IDSelectorAll : IDSelector {
 struct IDSelectorAnd : IDSelector {
     const IDSelector* lhs;
     const IDSelector* rhs;
-    IDSelectorAnd(const IDSelector* lhs, const IDSelector* rhs)
-            : lhs(lhs), rhs(rhs) {}
+    IDSelectorAnd(const IDSelector* lhs_, const IDSelector* rhs_)
+            : lhs(lhs_), rhs(rhs_) {}
     bool is_member(idx_t id) const final {
         return lhs->is_member(id) && rhs->is_member(id);
     }
@@ -149,8 +149,8 @@ struct IDSelectorAnd : IDSelector {
 struct IDSelectorOr : IDSelector {
     const IDSelector* lhs;
     const IDSelector* rhs;
-    IDSelectorOr(const IDSelector* lhs, const IDSelector* rhs)
-            : lhs(lhs), rhs(rhs) {}
+    IDSelectorOr(const IDSelector* lhs_, const IDSelector* rhs_)
+            : lhs(lhs_), rhs(rhs_) {}
     bool is_member(idx_t id) const final {
         return lhs->is_member(id) || rhs->is_member(id);
     }
@@ -162,8 +162,8 @@ struct IDSelectorOr : IDSelector {
 struct IDSelectorXOr : IDSelector {
     const IDSelector* lhs;
     const IDSelector* rhs;
-    IDSelectorXOr(const IDSelector* lhs, const IDSelector* rhs)
-            : lhs(lhs), rhs(rhs) {}
+    IDSelectorXOr(const IDSelector* lhs_, const IDSelector* rhs_)
+            : lhs(lhs_), rhs(rhs_) {}
     bool is_member(idx_t id) const final {
         return lhs->is_member(id) ^ rhs->is_member(id);
     }

data/vendor/faiss/faiss/impl/InvertedListScannerStats.h

@@ -0,0 +1,26 @@
+/*
+ * 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.
+ */
+
+#ifndef FAISS_INVERTED_LIST_SCANNER_STATS_H
+#define FAISS_INVERTED_LIST_SCANNER_STATS_H
+
+#include <cstddef>
+
+namespace faiss {
+
+/** Per-list statistics returned by inverted-list scanners. */
+struct InvertedListScannerStats {
+    /// Number of distances computed after IDSelector filtering.
+    size_t scan_cnt = 0;
+
+    /// Number of heap updates.
+    size_t nheap_updates = 0;
+};
+
+} // namespace faiss
+
+#endif

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

@@ -22,11 +22,11 @@
 #include <faiss/utils/distances.h>
 #include <faiss/utils/utils.h>
 
-#include <faiss/utils/approx_topk/approx_topk.h>
+#include <faiss/impl/approx_topk/approx_topk.h>
 
 // this is needed for prefetching
 
-#ifdef __AVX2__
+#ifdef COMPILE_SIMD_AVX2
 #include <xmmintrin.h>
 #endif
 
@@ -154,12 +154,15 @@ lsq::LSQTimer lsq_timer;
 using lsq::LSQTimerScope;
 
 LocalSearchQuantizer::LocalSearchQuantizer(
-        size_t d,
-        size_t M,
-        size_t nbits,
-        Search_type_t search_type)
-        : AdditiveQuantizer(d, std::vector<size_t>(M, nbits), search_type) {
-    K = (1 << nbits);
+        size_t d_in,
+        size_t M_in,
+        size_t nbits_in,
+        Search_type_t search_type_in)
+        : AdditiveQuantizer(
+                  d_in,
+                  std::vector<size_t>(M_in, nbits_in),
+                  search_type_in) {
+    K = (1 << nbits_in);
     std::srand(random_seed);
 }
 
@@ -170,7 +173,7 @@ LocalSearchQuantizer::~LocalSearchQuantizer() {
 LocalSearchQuantizer::LocalSearchQuantizer() : LocalSearchQuantizer(0, 0, 0) {}
 
 void LocalSearchQuantizer::train(size_t n, const float* x) {
-    FAISS_THROW_IF_NOT(K == (1 << nbits[0]));
+    FAISS_THROW_IF_NOT(K == static_cast<size_t>(1 << nbits[0]));
     nperts = std::min(nperts, M);
 
     lsq_timer.reset();
@@ -194,7 +197,7 @@ void LocalSearchQuantizer::train(size_t n, const float* x) {
     std::vector<float> stddev(d, 0);
 
 #pragma omp parallel for
-    for (int64_t i = 0; i < d; i++) {
+    for (int64_t i = 0; i < static_cast<int64_t>(d); i++) {
         float mean = 0;
         for (size_t j = 0; j < n; j++) {
             mean += x[j * d + i];
@@ -362,7 +365,7 @@ void LocalSearchQuantizer::update_codebooks(
         }
 
         // add a regularization term to B'B
-        for (int64_t i = 0; i < M * K; i++) {
+        for (size_t i = 0; i < M * K; i++) {
             bb[i * (M * K) + i] += lambd;
         }
 
@@ -427,7 +430,7 @@ void LocalSearchQuantizer::update_codebooks(
         }
 
         // add a regularization term to B'B
-        for (int64_t i = 0; i < M * K; i++) {
+        for (size_t i = 0; i < M * K; i++) {
             bb[i * (M * K) + i] += lambd;
         }
 
@@ -540,7 +543,7 @@ void LocalSearchQuantizer::icm_encode_impl(
         std::mt19937& gen,
         size_t n,
         size_t ils_iters,
-        bool verbose) const {
+        bool verbose_in) const {
     std::vector<float> unaries(n * M * K); // [M, n, K]
     compute_unary_terms(x, unaries.data(), n);
 
@@ -564,7 +567,7 @@ void LocalSearchQuantizer::icm_encode_impl(
 
         // select the best code for every vector xi
 #pragma omp parallel for reduction(+ : n_betters, mean_obj)
-        for (int64_t i = 0; i < n; i++) {
+        for (int64_t i = 0; i < static_cast<int64_t>(n); i++) {
             if (icm_objs[i] < best_objs[i]) {
                 best_objs[i] = icm_objs[i];
                 memcpy(best_codes.data() + i * M,
@@ -578,7 +581,7 @@ void LocalSearchQuantizer::icm_encode_impl(
 
         memcpy(codes, best_codes.data(), sizeof(int32_t) * n * M);
 
-        if (verbose) {
+        if (verbose_in) {
             printf("\tils_iter %zd: obj = %lf, n_betters/n = %zd/%zd\n",
                    iter1,
                    mean_obj,
@@ -597,73 +600,75 @@ void LocalSearchQuantizer::icm_encode_step(
     FAISS_THROW_IF_NOT(M != 0 && K != 0);
     FAISS_THROW_IF_NOT(binaries != nullptr);
 
+    // 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)
-    for (int64_t i = 0; i < n; i++) {
-        std::vector<float> objs(K);
-
-        for (size_t iter = 0; iter < n_iters; iter++) {
-            // condition on the m-th subcode
-            for (size_t m = 0; m < M; m++) {
-                // copy
-                auto u = unaries + m * n * K + i * K;
-                for (size_t code = 0; code < K; code++) {
-                    objs[code] = u[code];
-                }
-
-                // compute objective function by adding unary
-                // and binary terms together
-                for (size_t other_m = 0; other_m < M; other_m++) {
-                    if (other_m == m) {
-                        continue;
+        for (int64_t i = 0; i < static_cast<int64_t>(n); i++) {
+            std::vector<float> objs(K);
+
+            for (size_t iter = 0; iter < n_iters; iter++) {
+                // condition on the m-th subcode
+                for (size_t m = 0; m < M; m++) {
+                    // copy
+                    auto u = unaries + m * n * K + i * K;
+                    for (size_t code = 0; code < K; code++) {
+                        objs[code] = u[code];
                     }
 
-#ifdef __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);
+                    // compute objective function by adding unary
+                    // and binary terms together
+                    for (size_t other_m = 0; other_m < M; other_m++) {
+                        if (other_m == m) {
+                            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);
+                            }
                         }
-                    }
 #endif
 
-                    for (int32_t code = 0; code < K; code++) {
-                        int32_t code2 = codes[i * M + other_m];
-                        size_t binary_idx = other_m * M * K * K + m * K * K +
-                                code2 * K + code;
-                        // binaries[m, other_m, code, code2].
-                        // It is symmetric over (m <-> other_m)
-                        //   and (code <-> code2).
-                        // So, replace the op with
-                        //   binaries[other_m, m, code2, code].
-                        objs[code] += binaries[binary_idx];
+                        for (size_t code = 0; code < K; code++) {
+                            int32_t code2 = codes[i * M + other_m];
+                            size_t binary_idx = other_m * M * K * K +
+                                    m * K * K + code2 * K + code;
+                            // binaries[m, other_m, code, code2].
+                            // It is symmetric over (m <-> other_m)
+                            //   and (code <-> code2).
+                            // So, replace the op with
+                            //   binaries[other_m, m, code2, code].
+                            objs[code] += binaries[binary_idx];
+                        }
                     }
-                }
 
-                // find the optimal value of the m-th subcode
-                float best_obj = HUGE_VALF;
-                int32_t best_code = 0;
+                    // find the optimal value of the m-th subcode
+                    float best_obj = HUGE_VALF;
+                    int32_t best_code = 0;
 
-                // find one using SIMD. The following operation is similar
-                // to the search of the smallest element in objs
-                using C = CMax<float, int>;
-                HeapWithBuckets<C, 16, 1>::addn(
-                        K, objs.data(), 1, &best_obj, &best_code);
+                    // find one using SIMD. The following operation is similar
+                    // to the search of the smallest element in objs
+                    HeapWithBucketsCMaxFloat<16, 1, SL>::addn(
+                            K, objs.data(), 1, &best_obj, &best_code);
 
-                // done
-                codes[i * M + m] = best_code;
+                    // done
+                    codes[i * M + m] = best_code;
 
-            } // loop M
+                } // loop M
+            }
         }
-    }
+    });
 }
 void LocalSearchQuantizer::perturb_codes(
         int32_t* codes,
@@ -687,7 +692,7 @@ void LocalSearchQuantizer::compute_binary_terms(float* binaries) const {
 
     with_simd_level([&]<SIMDLevel SL>() {
 #pragma omp parallel for
-        for (int64_t m12 = 0; m12 < M * M; m12++) {
+        for (int64_t m12 = 0; m12 < static_cast<int64_t>(M * M); m12++) {
             size_t m1 = m12 / M;
             size_t m2 = m12 % M;
 
@@ -744,7 +749,7 @@ void LocalSearchQuantizer::compute_unary_terms(
     fvec_norms_L2sqr(norms.data(), codebooks.data(), d, M * K);
 
 #pragma omp parallel for
-    for (int64_t i = 0; i < n; i++) {
+    for (int64_t i = 0; i < static_cast<int64_t>(n); i++) {
         for (size_t m = 0; m < M; m++) {
             float* u = unaries + m * n * K + i * K;
             fvec_add(K, u, norms.data() + m * K, u);
@@ -766,7 +771,7 @@ float LocalSearchQuantizer::evaluate(
     with_simd_level([&]<SIMDLevel SL>() {
         float local_obj = 0.0f;
 #pragma omp parallel for reduction(+ : local_obj)
-        for (int64_t i = 0; i < n; i++) {
+        for (int64_t i = 0; i < static_cast<int64_t>(n); i++) {
             const auto code = codes + i * M;
             const auto decoded_i = decoded_x.data() + i * d;
             for (size_t m = 0; m < M; m++) {
@@ -791,8 +796,8 @@ float LocalSearchQuantizer::evaluate(
 
 namespace lsq {
 
-IcmEncoder::IcmEncoder(const LocalSearchQuantizer* lsq)
-        : verbose(false), lsq(lsq) {}
+IcmEncoder::IcmEncoder(const LocalSearchQuantizer* lsq_in)
+        : verbose(false), lsq(lsq_in) {}
 
 void IcmEncoder::set_binary_term() {
     auto M = lsq->M;
@@ -830,8 +835,8 @@ void LSQTimer::reset() {
     t.clear();
 }
 
-LSQTimerScope::LSQTimerScope(LSQTimer* timer, std::string name)
-        : timer(timer), name(std::move(name)), finished(false) {
+LSQTimerScope::LSQTimerScope(LSQTimer* timer_in, std::string name_in)
+        : timer(timer_in), name(std::move(name_in)), finished(false) {
     t0 = getmillisecs();
 }