faiss 0.3.0 → 0.3.2

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

@@ -5,8 +5,6 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-// -*- c++ -*-
-
 #pragma once
 
 #include <queue>
@@ -42,10 +40,13 @@ namespace faiss {
 struct VisitedTable;
 struct DistanceComputer; // from AuxIndexStructures
 struct HNSWStats;
+template <class C>
+struct ResultHandler;
 
 struct SearchParametersHNSW : SearchParameters {
     int efSearch = 16;
     bool check_relative_distance = true;
+    bool bounded_queue = true;
 
     ~SearchParametersHNSW() {}
 };
@@ -54,6 +55,9 @@ struct HNSW {
     /// internal storage of vectors (32 bits: this is expensive)
     using storage_idx_t = int32_t;
 
+    // for now we do only these distances
+    using C = CMax<float, int64_t>;
+
     typedef std::pair<float, storage_idx_t> Node;
 
     /** Heap structure that allows fast
@@ -138,9 +142,6 @@ struct HNSW {
     /// enough?
     bool check_relative_distance = true;
 
-    /// number of entry points in levels > 0.
-    int upper_beam = 1;
-
     /// use bounded queue during exploration
     bool search_bounded_queue = true;
 
@@ -181,7 +182,8 @@ struct HNSW {
             float d_nearest,
             int level,
             omp_lock_t* locks,
-            VisitedTable& vt);
+            VisitedTable& vt,
+            bool keep_max_size_level0 = false);
 
     /** add point pt_id on all levels <= pt_level and build the link
      * structure for them. */
@@ -190,29 +192,27 @@ struct HNSW {
             int pt_level,
             int pt_id,
             std::vector<omp_lock_t>& locks,
-            VisitedTable& vt);
+            VisitedTable& vt,
+            bool keep_max_size_level0 = false);
 
     /// search interface for 1 point, single thread
     HNSWStats search(
             DistanceComputer& qdis,
-            int k,
-            idx_t* I,
-            float* D,
+            ResultHandler<C>& res,
             VisitedTable& vt,
             const SearchParametersHNSW* params = nullptr) const;
 
     /// search only in level 0 from a given vertex
     void search_level_0(
             DistanceComputer& qdis,
-            int k,
-            idx_t* idxi,
-            float* simi,
+            ResultHandler<C>& res,
             idx_t nprobe,
             const storage_idx_t* nearest_i,
             const float* nearest_d,
             int search_type,
             HNSWStats& search_stats,
-            VisitedTable& vt) const;
+            VisitedTable& vt,
+            const SearchParametersHNSW* params = nullptr) const;
 
     void reset();
 
@@ -225,38 +225,60 @@ struct HNSW {
             DistanceComputer& qdis,
             std::priority_queue<NodeDistFarther>& input,
             std::vector<NodeDistFarther>& output,
-            int max_size);
+            int max_size,
+            bool keep_max_size_level0 = false);
+
+    void permute_entries(const idx_t* map);
 };
 
 struct HNSWStats {
-    size_t n1, n2, n3;
-    size_t ndis;
-    size_t nreorder;
-
-    HNSWStats(
-            size_t n1 = 0,
-            size_t n2 = 0,
-            size_t n3 = 0,
-            size_t ndis = 0,
-            size_t nreorder = 0)
-            : n1(n1), n2(n2), n3(n3), ndis(ndis), nreorder(nreorder) {}
+    size_t n1 = 0; /// number of vectors searched
+    size_t n2 =
+            0; /// number of queries for which the candidate list is exhausted
+    size_t ndis = 0;  /// number of distances computed
+    size_t nhops = 0; /// number of hops aka number of edges traversed
 
     void reset() {
-        n1 = n2 = n3 = 0;
+        n1 = n2 = 0;
         ndis = 0;
-        nreorder = 0;
+        nhops = 0;
     }
 
     void combine(const HNSWStats& other) {
         n1 += other.n1;
         n2 += other.n2;
-        n3 += other.n3;
         ndis += other.ndis;
-        nreorder += other.nreorder;
+        nhops += other.nhops;
     }
 };
 
 // global var that collects them all
 FAISS_API extern HNSWStats hnsw_stats;
 
+int search_from_candidates(
+        const HNSW& hnsw,
+        DistanceComputer& qdis,
+        ResultHandler<HNSW::C>& res,
+        HNSW::MinimaxHeap& candidates,
+        VisitedTable& vt,
+        HNSWStats& stats,
+        int level,
+        int nres_in = 0,
+        const SearchParametersHNSW* params = nullptr);
+
+HNSWStats greedy_update_nearest(
+        const HNSW& hnsw,
+        DistanceComputer& qdis,
+        int level,
+        HNSW::storage_idx_t& nearest,
+        float& d_nearest);
+
+std::priority_queue<HNSW::Node> search_from_candidate_unbounded(
+        const HNSW& hnsw,
+        const HNSW::Node& node,
+        DistanceComputer& qdis,
+        int ef,
+        VisitedTable* vt,
+        HNSWStats& stats);
+
 } // namespace faiss

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

@@ -10,7 +10,7 @@
 #include <unordered_set>
 #include <vector>
 
-#include <faiss/Index.h>
+#include <faiss/MetricType.h>
 
 /** IDSelector is intended to define a subset of vectors to handle (for removal
  * or as subset to search) */
@@ -140,7 +140,7 @@ struct IDSelectorAnd : IDSelector {
             : lhs(lhs), rhs(rhs) {}
     bool is_member(idx_t id) const final {
         return lhs->is_member(id) && rhs->is_member(id);
-    };
+    }
     virtual ~IDSelectorAnd() {}
 };
 
@@ -153,7 +153,7 @@ struct IDSelectorOr : IDSelector {
             : lhs(lhs), rhs(rhs) {}
     bool is_member(idx_t id) const final {
         return lhs->is_member(id) || rhs->is_member(id);
-    };
+    }
     virtual ~IDSelectorOr() {}
 };
 
@@ -166,7 +166,7 @@ struct IDSelectorXOr : IDSelector {
             : lhs(lhs), rhs(rhs) {}
     bool is_member(idx_t id) const final {
         return lhs->is_member(id) ^ rhs->is_member(id);
-    };
+    }
     virtual ~IDSelectorXOr() {}
 };
 

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

@@ -104,10 +104,10 @@ int dgemm_(
 
 namespace {
 
-void fmat_inverse(float* a, int n) {
-    int info;
-    int lwork = n * n;
-    std::vector<int> ipiv(n);
+void fmat_inverse(float* a, FINTEGER n) {
+    FINTEGER info;
+    FINTEGER lwork = n * n;
+    std::vector<FINTEGER> ipiv(n);
     std::vector<float> workspace(lwork);
 
     sgetrf_(&n, &n, a, &n, ipiv.data(), &info);
@@ -123,10 +123,10 @@ void dfvec_add(size_t d, const double* a, const float* b, double* c) {
     }
 }
 
-void dmat_inverse(double* a, int n) {
-    int info;
-    int lwork = n * n;
-    std::vector<int> ipiv(n);
+void dmat_inverse(double* a, FINTEGER n) {
+    FINTEGER info;
+    FINTEGER lwork = n * n;
+    std::vector<FINTEGER> ipiv(n);
     std::vector<double> workspace(lwork);
 
     dgetrf_(&n, &n, a, &n, ipiv.data(), &info);
@@ -628,7 +628,9 @@ void LocalSearchQuantizer::icm_encode_step(
                         {
                             size_t binary_idx = (other_m + 1) * M * K * K +
                                     m * K * K + code2 * K + code;
-                            _mm_prefetch(binaries + binary_idx, _MM_HINT_T0);
+                            _mm_prefetch(
+                                    (const char*)(binaries + binary_idx),
+                                    _MM_HINT_T0);
                         }
                     }
 #endif

data/vendor/faiss/faiss/impl/LookupTableScaler.h

@@ -38,6 +38,23 @@ struct DummyScaler {
         return simd16uint16(0);
     }
 
+#ifdef __AVX512F__
+    inline simd64uint8 lookup(const simd64uint8&, const simd64uint8&) const {
+        FAISS_THROW_MSG("DummyScaler::lookup should not be called.");
+        return simd64uint8(0);
+    }
+
+    inline simd32uint16 scale_lo(const simd64uint8&) const {
+        FAISS_THROW_MSG("DummyScaler::scale_lo should not be called.");
+        return simd32uint16(0);
+    }
+
+    inline simd32uint16 scale_hi(const simd64uint8&) const {
+        FAISS_THROW_MSG("DummyScaler::scale_hi should not be called.");
+        return simd32uint16(0);
+    }
+#endif
+
     template <class dist_t>
     inline dist_t scale_one(const dist_t&) const {
         FAISS_THROW_MSG("DummyScaler::scale_one should not be called.");
@@ -67,6 +84,23 @@ struct NormTableScaler {
         return (simd16uint16(res) >> 8) * scale_simd;
     }
 
+#ifdef __AVX512F__
+    inline simd64uint8 lookup(const simd64uint8& lut, const simd64uint8& c)
+            const {
+        return lut.lookup_4_lanes(c);
+    }
+
+    inline simd32uint16 scale_lo(const simd64uint8& res) const {
+        auto scale_simd_wide = simd32uint16(scale_simd, scale_simd);
+        return simd32uint16(res) * scale_simd_wide;
+    }
+
+    inline simd32uint16 scale_hi(const simd64uint8& res) const {
+        auto scale_simd_wide = simd32uint16(scale_simd, scale_simd);
+        return (simd32uint16(res) >> 8) * scale_simd_wide;
+    }
+#endif
+
     // for non-SIMD implem 2, 3, 4
     template <class dist_t>
     inline dist_t scale_one(const dist_t& x) const {

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

@@ -154,15 +154,20 @@ NNDescent::NNDescent(const int d, const int K) : K(K), d(d) {
 NNDescent::~NNDescent() {}
 
 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)
-    for (int n = 0; n < ntotal; n++) {
-        graph[n].join([&](int i, int j) {
-            if (i != j) {
-                float dist = qdis.symmetric_dis(i, j);
-                graph[i].insert(j, dist);
-                graph[j].insert(i, dist);
-            }
-        });
+        for (idx_t n = i0; n < i1; n++) {
+            graph[n].join([&](int i, int j) {
+                if (i != j) {
+                    float dist = qdis.symmetric_dis(i, j);
+                    graph[i].insert(j, dist);
+                    graph[j].insert(i, dist);
+                }
+            });
+        }
+        InterruptCallback::check();
     }
 }
 
@@ -195,8 +200,9 @@ void NNDescent::update() {
         int l = 0;
 
         while ((l < maxl) && (c < S)) {
-            if (nn.pool[l].flag)
+            if (nn.pool[l].flag) {
                 ++c;
+            }
             ++l;
         }
         nn.M = l;
@@ -305,8 +311,9 @@ void NNDescent::generate_eval_set(
     for (int i = 0; i < c.size(); i++) {
         std::vector<Neighbor> tmp;
         for (int j = 0; j < N; j++) {
-            if (c[i] == j)
+            if (c[i] == j) {
                 continue; // skip itself
+            }
             float dist = qdis.symmetric_dis(c[i], j);
             tmp.push_back(Neighbor(j, dist, true));
         }
@@ -360,8 +367,9 @@ void NNDescent::init_graph(DistanceComputer& qdis) {
 
             for (int j = 0; j < S; j++) {
                 int id = tmp[j];
-                if (id == i)
+                if (id == i) {
                     continue;
+                }
                 float dist = qdis.symmetric_dis(i, id);
 
                 graph[i].pool.push_back(Neighbor(id, dist, true));
@@ -374,6 +382,10 @@ void NNDescent::init_graph(DistanceComputer& qdis) {
 
 void NNDescent::build(DistanceComputer& qdis, const int n, bool verbose) {
     FAISS_THROW_IF_NOT_MSG(L >= K, "L should be >= K in NNDescent.build");
+    FAISS_THROW_IF_NOT_FMT(
+            n > NUM_EVAL_POINTS,
+            "NNDescent.build cannot build a graph smaller than %d",
+            int(NUM_EVAL_POINTS));
 
     if (verbose) {
         printf("Parameters: K=%d, S=%d, R=%d, L=%d, iter=%d\n",
@@ -403,7 +415,7 @@ void NNDescent::build(DistanceComputer& qdis, const int n, bool verbose) {
     has_built = true;
 
     if (verbose) {
-        printf("Addes %d points into the index\n", ntotal);
+        printf("Added %d points into the index\n", ntotal);
     }
 }
 
@@ -414,30 +426,30 @@ void NNDescent::search(
         float* dists,
         VisitedTable& vt) const {
     FAISS_THROW_IF_NOT_MSG(has_built, "The index is not build yet.");
-    int L = std::max(search_L, topk);
+    int L_2 = std::max(search_L, topk);
 
     // candidate pool, the K best items is the result.
-    std::vector<Neighbor> retset(L + 1);
+    std::vector<Neighbor> retset(L_2 + 1);
 
-    // Randomly choose L points to initialize the candidate pool
-    std::vector<int> init_ids(L);
+    // Randomly choose L_2 points to initialize the candidate pool
+    std::vector<int> init_ids(L_2);
     std::mt19937 rng(random_seed);
 
-    gen_random(rng, init_ids.data(), L, ntotal);
-    for (int i = 0; i < L; i++) {
+    gen_random(rng, init_ids.data(), L_2, ntotal);
+    for (int i = 0; i < L_2; i++) {
         int id = init_ids[i];
         float dist = qdis(id);
         retset[i] = Neighbor(id, dist, true);
     }
 
     // Maintain the candidate pool in ascending order
-    std::sort(retset.begin(), retset.begin() + L);
+    std::sort(retset.begin(), retset.begin() + L_2);
 
     int k = 0;
 
-    // Stop until the smallest position updated is >= L
-    while (k < L) {
-        int nk = L;
+    // Stop until the smallest position updated is >= L_2
+    while (k < L_2) {
+        int nk = L_2;
 
         if (retset[k].flag) {
             retset[k].flag = false;
@@ -445,25 +457,28 @@ void NNDescent::search(
 
             for (int m = 0; m < K; ++m) {
                 int id = final_graph[n * K + m];
-                if (vt.get(id))
+                if (vt.get(id)) {
                     continue;
+                }
 
                 vt.set(id);
                 float dist = qdis(id);
-                if (dist >= retset[L - 1].distance)
+                if (dist >= retset[L_2 - 1].distance) {
                     continue;
+                }
 
                 Neighbor nn(id, dist, true);
-                int r = insert_into_pool(retset.data(), L, nn);
+                int r = insert_into_pool(retset.data(), L_2, nn);
 
                 if (r < nk)
                     nk = r;
             }
         }
-        if (nk <= k)
+        if (nk <= k) {
             k = nk;
-        else
+        } else {
             ++k;
+        }
     }
     for (size_t i = 0; i < topk; i++) {
         indices[i] = retset[i].id;

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

@@ -25,35 +25,6 @@ namespace {
 // It needs to be smaller than 0
 constexpr int EMPTY_ID = -1;
 
-/* Wrap the distance computer into one that negates the
-   distances. This makes supporting INNER_PRODUCE search easier */
-
-struct NegativeDistanceComputer : DistanceComputer {
-    /// owned by this
-    DistanceComputer* basedis;
-
-    explicit NegativeDistanceComputer(DistanceComputer* basedis)
-            : basedis(basedis) {}
-
-    void set_query(const float* x) override {
-        basedis->set_query(x);
-    }
-
-    /// compute distance of vector i to current query
-    float operator()(idx_t i) override {
-        return -(*basedis)(i);
-    }
-
-    /// compute distance between two stored vectors
-    float symmetric_dis(idx_t i, idx_t j) override {
-        return -basedis->symmetric_dis(i, j);
-    }
-
-    ~NegativeDistanceComputer() override {
-        delete basedis;
-    }
-};
-
 } // namespace
 
 DistanceComputer* storage_distance_computer(const Index* storage) {

data/vendor/faiss/faiss/impl/NSG.h

@@ -54,7 +54,7 @@ namespace nsg {
 
 template <class node_t>
 struct Graph {
-    node_t* data;    ///< the flattened adjacency matrix
+    node_t* data;    ///< the flattened adjacency matrix, size N-by-K
     int K;           ///< nb of neighbors per node
     int N;           ///< total nb of nodes
     bool own_fields; ///< the underlying data owned by itself or not

data/vendor/faiss/faiss/impl/PolysemousTraining.cpp

@@ -12,11 +12,11 @@
 #include <omp.h>
 #include <stdint.h>
 
+#include <algorithm>
 #include <cmath>
 #include <cstdlib>
 #include <cstring>
-
-#include <algorithm>
+#include <memory>
 
 #include <faiss/utils/distances.h>
 #include <faiss/utils/hamming.h>
@@ -683,18 +683,21 @@ struct RankingScore2 : Score3Computer<float, double> {
     double accum_gt_weight_diff(
             const std::vector<int>& a,
             const std::vector<int>& b) {
-        int nb = b.size(), na = a.size();
+        const auto nb_2 = b.size();
+        const auto na = a.size();
 
         double accu = 0;
-        int j = 0;
-        for (int i = 0; i < na; i++) {
-            int ai = a[i];
-            while (j < nb && ai >= b[j])
+        size_t j = 0;
+        for (size_t i = 0; i < na; i++) {
+            const auto ai = a[i];
+            while (j < nb_2 && ai >= b[j]) {
                 j++;
+            }
 
             double accu_i = 0;
-            for (int k = j; k < b.size(); k++)
+            for (auto k = j; k < b.size(); k++) {
                 accu_i += rank_weight(b[k] - ai);
+            }
 
             accu += rank_weight(ai) * accu_i;
         }
@@ -882,14 +885,13 @@ void PolysemousTraining::optimize_ranking(
 
         double t0 = getmillisecs();
 
-        PermutationObjective* obj = new RankingScore2(
+        std::unique_ptr<PermutationObjective> obj(new RankingScore2(
                 nbits,
                 nq,
                 nb,
                 codes.data(),
                 codes.data() + nq,
-                gt_distances.data());
-        ScopeDeleter1<PermutationObjective> del(obj);
+                gt_distances.data()));
 
         if (verbose > 0) {
             printf("   m=%d, nq=%zd, nb=%zd, initialize RankingScore "
@@ -900,7 +902,7 @@ void PolysemousTraining::optimize_ranking(
                    getmillisecs() - t0);
         }
 
-        SimulatedAnnealingOptimizer optim(obj, *this);
+        SimulatedAnnealingOptimizer optim(obj.get(), *this);
 
         if (log_pattern.size()) {
             char fname[256];

data/vendor/faiss/faiss/impl/ProductAdditiveQuantizer.h

@@ -151,4 +151,4 @@ struct ProductResidualQuantizer : ProductAdditiveQuantizer {
     ProductResidualQuantizer();
 };
 
-}; // namespace faiss
+} // namespace faiss

data/vendor/faiss/faiss/impl/ProductQuantizer.cpp

@@ -61,6 +61,7 @@ void ProductQuantizer::set_derived_values() {
             "The dimension of the vector (d) should be a multiple of the number of subquantizers (M)");
     dsub = d / M;
     code_size = (nbits * M + 7) / 8;
+    FAISS_THROW_IF_MSG(nbits > 24, "nbits larger than 24 is not practical.");
     ksub = 1 << nbits;
     centroids.resize(d * ksub);
     verbose = false;
@@ -135,11 +136,10 @@ void ProductQuantizer::train(size_t n, const float* x) {
             }
         }
 
-        float* xslice = new float[n * dsub];
-        ScopeDeleter<float> del(xslice);
+        std::unique_ptr<float[]> xslice(new float[n * dsub]);
         for (int m = 0; m < M; m++) {
             for (int j = 0; j < n; j++)
-                memcpy(xslice + j * dsub,
+                memcpy(xslice.get() + j * dsub,
                        x + j * d + m * dsub,
                        dsub * sizeof(float));
 
@@ -153,11 +153,19 @@ void ProductQuantizer::train(size_t n, const float* x) {
             switch (final_train_type) {
                 case Train_hypercube:
                     init_hypercube(
-                            dsub, nbits, n, xslice, clus.centroids.data());
+                            dsub,
+                            nbits,
+                            n,
+                            xslice.get(),
+                            clus.centroids.data());
                     break;
                 case Train_hypercube_pca:
                     init_hypercube_pca(
-                            dsub, nbits, n, xslice, clus.centroids.data());
+                            dsub,
+                            nbits,
+                            n,
+                            xslice.get(),
+                            clus.centroids.data());
                     break;
                 case Train_hot_start:
                     memcpy(clus.centroids.data(),
@@ -172,7 +180,7 @@ void ProductQuantizer::train(size_t n, const float* x) {
                 printf("Training PQ slice %d/%zd\n", m, M);
             }
             IndexFlatL2 index(dsub);
-            clus.train(n, xslice, assign_index ? *assign_index : index);
+            clus.train(n, xslice.get(), assign_index ? *assign_index : index);
             set_params(clus.centroids.data(), m);
         }
 
@@ -306,7 +314,8 @@ void ProductQuantizer::decode(const uint8_t* code, float* x) const {
 }
 
 void ProductQuantizer::decode(const uint8_t* code, float* x, size_t n) const {
-    for (size_t i = 0; i < n; i++) {
+#pragma omp parallel for if (n > 100)
+    for (int64_t i = 0; i < n; i++) {
         this->decode(code + code_size * i, x + d * i);
     }
 }
@@ -342,21 +351,20 @@ void ProductQuantizer::compute_codes_with_assign_index(
         assign_index->reset();
         assign_index->add(ksub, get_centroids(m, 0));
         size_t bs = 65536;
-        float* xslice = new float[bs * dsub];
-        ScopeDeleter<float> del(xslice);
-        idx_t* assign = new idx_t[bs];
-        ScopeDeleter<idx_t> del2(assign);
+
+        std::unique_ptr<float[]> xslice(new float[bs * dsub]);
+        std::unique_ptr<idx_t[]> assign(new idx_t[bs]);
 
         for (size_t i0 = 0; i0 < n; i0 += bs) {
             size_t i1 = std::min(i0 + bs, n);
 
             for (size_t i = i0; i < i1; i++) {
-                memcpy(xslice + (i - i0) * dsub,
+                memcpy(xslice.get() + (i - i0) * dsub,
                        x + i * d + m * dsub,
                        dsub * sizeof(float));
             }
 
-            assign_index->assign(i1 - i0, xslice, assign);
+            assign_index->assign(i1 - i0, xslice.get(), assign.get());
 
             if (nbits == 8) {
                 uint8_t* c = codes + code_size * i0 + m;
@@ -405,15 +413,14 @@ void ProductQuantizer::compute_codes(const float* x, uint8_t* codes, size_t n)
         for (int64_t i = 0; i < n; i++)
             compute_code(x + i * d, codes + i * code_size);
 
-    } else { // worthwile to use BLAS
-        float* dis_tables = new float[n * ksub * M];
-        ScopeDeleter<float> del(dis_tables);
-        compute_distance_tables(n, x, dis_tables);
+    } else { // worthwhile to use BLAS
+        std::unique_ptr<float[]> dis_tables(new float[n * ksub * M]);
+        compute_distance_tables(n, x, dis_tables.get());
 
 #pragma omp parallel for
         for (int64_t i = 0; i < n; i++) {
             uint8_t* code = codes + i * code_size;
-            const float* tab = dis_tables + i * ksub * M;
+            const float* tab = dis_tables.get() + i * ksub * M;
             compute_code_from_distance_table(tab, code);
         }
     }
@@ -774,10 +781,6 @@ void ProductQuantizer::search_ip(
             init_finalize_heap);
 }
 
-static float sqr(float x) {
-    return x * x;
-}
-
 void ProductQuantizer::compute_sdc_table() {
     sdc_table.resize(M * ksub * ksub);