faiss 0.6.0 → 0.6.2

data/vendor/faiss/faiss/impl/ResidualQuantizer.cpp

@@ -70,23 +70,26 @@ ResidualQuantizer::ResidualQuantizer() {
 }
 
 ResidualQuantizer::ResidualQuantizer(
-        size_t d,
-        const std::vector<size_t>& nbits,
-        Search_type_t search_type)
+        size_t d_in,
+        const std::vector<size_t>& nbits_in,
+        Search_type_t search_type_in)
         : ResidualQuantizer() {
-    this->search_type = search_type;
-    this->d = d;
-    M = nbits.size();
-    this->nbits = nbits;
+    this->search_type = search_type_in;
+    this->d = d_in;
+    M = nbits_in.size();
+    this->nbits = nbits_in;
     set_derived_values();
 }
 
 ResidualQuantizer::ResidualQuantizer(
-        size_t d,
-        size_t M,
-        size_t nbits,
-        Search_type_t search_type)
-        : ResidualQuantizer(d, std::vector<size_t>(M, nbits), search_type) {}
+        size_t d_in,
+        size_t M_in,
+        size_t nbits_in,
+        Search_type_t search_type_in)
+        : ResidualQuantizer(
+                  d_in,
+                  std::vector<size_t>(M_in, nbits_in),
+                  search_type_in) {}
 
 void ResidualQuantizer::initialize_from(
         const ResidualQuantizer& other,
@@ -142,7 +145,7 @@ void ResidualQuantizer::train(size_t n, const float* x) {
     double t0 = getmillisecs();
     double clustering_time = 0;
 
-    for (int m = 0; m < M; m++) {
+    for (size_t m = 0; m < M; m++) {
         int K = 1 << nbits[m];
 
         // on which residuals to train
@@ -157,7 +160,7 @@ void ResidualQuantizer::train(size_t n, const float* x) {
             }
             train_residuals = residuals1;
         }
-        std::vector<float> codebooks;
+        std::vector<float> cur_codebooks;
         float obj = 0;
 
         std::unique_ptr<Index> assign_index;
@@ -175,7 +178,7 @@ void ResidualQuantizer::train(size_t n, const float* x) {
                     train_residuals.size() / d,
                     train_residuals.data(),
                     *assign_index.get());
-            codebooks.swap(clus.centroids);
+            cur_codebooks.swap(clus.centroids);
             assign_index->reset();
             obj = clus.iteration_stats.back().obj;
         } else { // progressive dim clustering
@@ -185,14 +188,14 @@ void ResidualQuantizer::train(size_t n, const float* x) {
                     train_residuals.size() / d,
                     train_residuals.data(),
                     assign_index_factory ? *assign_index_factory : default_fac);
-            codebooks.swap(clus.centroids);
+            cur_codebooks.swap(clus.centroids);
             obj = clus.iteration_stats.back().obj;
         }
         clustering_time += (getmillisecs() - t1) / 1000;
 
         memcpy(this->codebooks.data() + codebook_offsets[m] * d,
-               codebooks.data(),
-               codebooks.size() * sizeof(codebooks[0]));
+               cur_codebooks.data(),
+               cur_codebooks.size() * sizeof(cur_codebooks[0]));
 
         // quantize using the new codebooks
 
@@ -221,7 +224,7 @@ void ResidualQuantizer::train(size_t n, const float* x) {
             beam_search_encode_step(
                     d,
                     K,
-                    codebooks.data(),
+                    cur_codebooks.data(),
                     i1 - i0,
                     cur_beam_size,
                     residuals.data() + i0 * cur_beam_size * d,
@@ -239,16 +242,16 @@ void ResidualQuantizer::train(size_t n, const float* x) {
         distances.swap(new_distances);
 
         float sum_distances = 0;
-        for (int j = 0; j < distances.size(); j++) {
+        for (size_t j = 0; j < distances.size(); j++) {
             sum_distances += distances[j];
         }
 
         if (verbose) {
-            printf("[%.3f s, %.3f s clustering] train stage %d, %d bits, kmeans objective %g, "
+            printf("[%.3f s, %.3f s clustering] train stage %zd, %d bits, kmeans objective %g, "
                    "total distance %g, beam_size %d->%d (batch size %zd)\n",
                    (getmillisecs() - t0) / 1000,
                    clustering_time,
-                   m,
+                   size_t(m),
                    int(nbits[m]),
                    obj,
                    sum_distances,
@@ -314,7 +317,7 @@ float ResidualQuantizer::retrain_AQ_codebook(size_t n, const float* x) {
     std::vector<float> C(n * total_codebook_size);
     for (size_t i = 0; i < n; i++) {
         BitstringReader bsr(codes.data() + i * code_size, code_size);
-        for (int m = 0; m < M; m++) {
+        for (size_t m = 0; m < M; m++) {
             int idx = bsr.read(nbits[m]);
             C[i + (codebook_offsets[m] + idx) * n] = 1;
         }

data/vendor/faiss/faiss/impl/ResidualQuantizer.h

@@ -13,7 +13,7 @@
 #include <faiss/Clustering.h>
 #include <faiss/impl/AdditiveQuantizer.h>
 
-#include <faiss/utils/approx_topk/mode.h>
+#include <faiss/impl/approx_topk/approx_topk.h>
 
 namespace faiss {
 

data/vendor/faiss/faiss/impl/ResultHandler.h

@@ -15,6 +15,7 @@
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/FaissException.h>
 #include <faiss/impl/IDSelector.h>
+#include <faiss/impl/InvertedListScannerStats.h>
 #include <faiss/utils/Heap.h>
 #include <faiss/utils/partitioning.h>
 #include <algorithm>
@@ -29,6 +30,9 @@ struct ResultHandlerUnordered {
     // if not better than threshold, then not necessary to call add_result
     T threshold{};
 
+    // per-list scan statistics populated by inverted-list scanners
+    InvertedListScannerStats stats;
+
     // return whether threshold was updated
     virtual bool add_result(T dis, TI idx) = 0;
 
@@ -66,8 +70,8 @@ struct BlockResultHandler {
     size_t nq; // number of queries for which we search
     const IDSelector* sel;
 
-    explicit BlockResultHandler(size_t nq, const IDSelector* sel = nullptr)
-            : nq(nq), sel(sel) {
+    explicit BlockResultHandler(size_t nq_, const IDSelector* sel_ = nullptr)
+            : nq(nq_), sel(sel_) {
         assert(!use_sel || sel);
     }
 
@@ -108,10 +112,10 @@ struct SingleQueryBlockResultHandler : BlockResultHandler<C, use_sel> {
     ResultHandler& the_handler;
 
     explicit SingleQueryBlockResultHandler(
-            ResultHandler& the_handler,
-            const IDSelector* sel = nullptr)
-            : BlockResultHandler<C, use_sel>(1, sel),
-              the_handler(the_handler) {}
+            ResultHandler& the_handler_in,
+            const IDSelector* sel_in = nullptr)
+            : BlockResultHandler<C, use_sel>(1, sel_in),
+              the_handler(the_handler_in) {}
 
     struct SingleResultHandler : ResultHandlerT<C> {
         ResultHandler& the_handler;
@@ -121,9 +125,7 @@ struct SingleQueryBlockResultHandler : BlockResultHandler<C, use_sel> {
                 : the_handler(hr.the_handler) {}
 
         /// begin results for query # i
-        void begin(const size_t qid) {
-            assert(qid == 0);
-        }
+        void begin(const size_t /* qid */) {}
 
         /// add one result for query i
         bool add_result(T dis, TI idx) final {
@@ -149,15 +151,15 @@ struct TopkBlockResultHandler : BlockResultHandler<C, use_sel> {
     int64_t k; // number of results to keep
 
     TopkBlockResultHandler(
-            size_t nq,
-            T* dis_tab,
-            TI* ids_tab,
-            size_t k,
-            const IDSelector* sel = nullptr)
-            : BlockResultHandler<C, use_sel>(nq, sel),
-              dis_tab(dis_tab),
-              ids_tab(ids_tab),
-              k(k) {}
+            size_t nq_,
+            T* dis_tab_,
+            TI* ids_tab_,
+            size_t k_,
+            const IDSelector* sel_ = nullptr)
+            : BlockResultHandler<C, use_sel>(nq_, sel_),
+              dis_tab(dis_tab_),
+              ids_tab(ids_tab_),
+              k(k_) {}
 
     ~TopkBlockResultHandler() {}
 };
@@ -176,12 +178,16 @@ struct Top1BlockResultHandler : TopkBlockResultHandler<C, use_sel> {
     using BlockResultHandler<C, use_sel>::i1;
 
     Top1BlockResultHandler(
-            size_t nq,
-            T* dis_tab,
-            TI* ids_tab,
-            const IDSelector* sel = nullptr)
-            : TopkBlockResultHandler<C, use_sel>(nq, dis_tab, ids_tab, 1, sel) {
-    }
+            size_t nq_,
+            T* dis_tab_,
+            TI* ids_tab_,
+            const IDSelector* sel_ = nullptr)
+            : TopkBlockResultHandler<C, use_sel>(
+                      nq_,
+                      dis_tab_,
+                      ids_tab_,
+                      1,
+                      sel_) {}
 
     struct SingleResultHandler : ResultHandlerT<C> {
         Top1BlockResultHandler& hr;
@@ -190,7 +196,7 @@ struct Top1BlockResultHandler : TopkBlockResultHandler<C, use_sel> {
         TI min_idx;
         size_t current_idx = 0;
 
-        explicit SingleResultHandler(Top1BlockResultHandler& hr) : hr(hr) {}
+        explicit SingleResultHandler(Top1BlockResultHandler& hr_) : hr(hr_) {}
 
         /// begin results for query # i
         void begin(const size_t current_idx_2) {
@@ -217,18 +223,19 @@ struct Top1BlockResultHandler : TopkBlockResultHandler<C, use_sel> {
     };
 
     /// begin
-    void begin_multiple(size_t i0, size_t i1) final {
-        this->i0 = i0;
-        this->i1 = i1;
+    void begin_multiple(size_t i0_in, size_t i1_in) final {
+        this->i0 = i0_in;
+        this->i1 = i1_in;
 
         for (size_t i = i0; i < i1; i++) {
             this->dis_tab[i] = C::neutral();
+            this->ids_tab[i] = -1;
         }
     }
 
     /// add results for query i0..i1 and j0..j1
     void add_results(size_t j0, size_t j1, const T* dis_tab_2) final {
-        for (int64_t i = i0; i < i1; i++) {
+        for (size_t i = i0; i < i1; i++) {
             const T* dis_tab_i = dis_tab_2 + (j1 - j0) * (i - i0) - j0;
 
             auto& min_distance = this->dis_tab[i];
@@ -270,8 +277,8 @@ struct HeapResultHandler : ResultHandlerT<C> {
     T* heap_dis;
     TI* heap_ids;
 
-    HeapResultHandler(size_t k, T* heap_dis, TI* heap_ids)
-            : k(k), heap_dis(heap_dis), heap_ids(heap_ids) {
+    HeapResultHandler(size_t k_, T* heap_dis_, TI* heap_ids_)
+            : k(k_), heap_dis(heap_dis_), heap_ids(heap_ids_) {
         if (heap_dis) {
             this->threshold = heap_dis[0];
         }
@@ -297,13 +304,17 @@ struct HeapBlockResultHandler : TopkBlockResultHandler<C, use_sel> {
     using TopkBlockResultHandler<C, use_sel>::k;
 
     HeapBlockResultHandler(
-            size_t nq,
-            T* dis_tab,
-            TI* ids_tab,
-            size_t k,
-            const IDSelector* sel = nullptr)
-            : TopkBlockResultHandler<C, use_sel>(nq, dis_tab, ids_tab, k, sel) {
-    }
+            size_t nq_,
+            T* dis_tab_,
+            TI* ids_tab_,
+            size_t k_,
+            const IDSelector* sel_ = nullptr)
+            : TopkBlockResultHandler<C, use_sel>(
+                      nq_,
+                      dis_tab_,
+                      ids_tab_,
+                      k_,
+                      sel_) {}
 
     /******************************************************
      * API for 1 result at a time (each SingleResultHandler is
@@ -313,9 +324,9 @@ struct HeapBlockResultHandler : TopkBlockResultHandler<C, use_sel> {
     struct SingleResultHandler : HeapResultHandler<C, use_sel> {
         HeapBlockResultHandler& hr;
 
-        explicit SingleResultHandler(HeapBlockResultHandler& hr)
-                : HeapResultHandler<C, use_sel>(hr.k, nullptr, nullptr),
-                  hr(hr) {}
+        explicit SingleResultHandler(HeapBlockResultHandler& hr_)
+                : HeapResultHandler<C, use_sel>(hr_.k, nullptr, nullptr),
+                  hr(hr_) {}
 
         /// begin results for query # i
         void begin(size_t i) {
@@ -346,12 +357,14 @@ struct HeapBlockResultHandler : TopkBlockResultHandler<C, use_sel> {
     }
 
     /// add results for query i0..i1 and j0..j1
-    void add_results(size_t j0, size_t j1, const T* dis_tab) final {
+    void add_results(size_t j0, size_t j1, const T* dis_in) final {
 #pragma omp parallel for
-        for (int64_t i = i0; i < i1; i++) {
+        for (int64_t i = static_cast<int64_t>(i0); i < static_cast<int64_t>(i1);
+             i++) {
             T* heap_dis = this->dis_tab + i * k;
             TI* heap_ids = this->ids_tab + i * k;
-            const T* dis_tab_i = dis_tab + (j1 - j0) * (i - i0) - j0;
+            const T* dis_tab_i =
+                    dis_in + (j1 - j0) * (i - static_cast<int64_t>(i0)) - j0;
             T thresh = heap_dis[0];
             for (size_t j = j0; j < j1; j++) {
                 T dis = dis_tab_i[j];
@@ -365,8 +378,9 @@ struct HeapBlockResultHandler : TopkBlockResultHandler<C, use_sel> {
 
     /// series of results for queries i0..i1 is done
     void end_multiple() final {
-        // maybe parallel for
-        for (size_t i = i0; i < i1; i++) {
+#pragma omp parallel for schedule(static) if ((i1 - i0) * k >= 1024)
+        for (int64_t i = static_cast<int64_t>(i0); i < static_cast<int64_t>(i1);
+             i++) {
             heap_reorder<C>(k, this->dis_tab + i * k, this->ids_tab + i * k);
         }
     }
@@ -397,8 +411,8 @@ struct ReservoirTopN : ResultHandlerT<C> {
 
     ReservoirTopN() {}
 
-    ReservoirTopN(size_t n, size_t capacity, T* vals, TI* ids)
-            : vals(vals), ids(ids), i(0), n(n), capacity(capacity) {
+    ReservoirTopN(size_t n_, size_t capacity_, T* vals_, TI* ids_)
+            : vals(vals_), ids(ids_), i(0), n(n_), capacity(capacity_) {
         assert(n < capacity);
         threshold = C::neutral();
     }
@@ -436,7 +450,7 @@ struct ReservoirTopN : ResultHandlerT<C> {
     }
 
     void to_result(T* heap_dis, TI* heap_ids) const {
-        for (int j = 0; j < std::min(i, n); j++) {
+        for (size_t j = 0; j < std::min(i, n); j++) {
             heap_push<C>(j + 1, heap_dis, heap_ids, vals[j], ids[j]);
         }
 
@@ -458,16 +472,22 @@ struct ReservoirBlockResultHandler : TopkBlockResultHandler<C, use_sel> {
     using TI = typename C::TI;
     using BlockResultHandler<C, use_sel>::i0;
     using BlockResultHandler<C, use_sel>::i1;
+    using TopkBlockResultHandler<C, use_sel>::k;
 
     size_t capacity; // capacity of the reservoirs
 
     ReservoirBlockResultHandler(
-            size_t nq,
-            T* dis_tab,
-            TI* ids_tab,
-            size_t k,
-            const IDSelector* sel = nullptr)
-            : TopkBlockResultHandler<C, use_sel>(nq, dis_tab, ids_tab, k, sel) {
+            size_t nq_,
+            T* dis_tab_,
+            TI* ids_tab_,
+            size_t k_,
+            const IDSelector* sel_ = nullptr)
+            : TopkBlockResultHandler<C, use_sel>(
+                      nq_,
+                      dis_tab_,
+                      ids_tab_,
+                      k_,
+                      sel_) {
         // double then round up to multiple of 16 (for SIMD alignment)
         capacity = (2 * k + 15) & ~15;
     }
@@ -483,9 +503,9 @@ struct ReservoirBlockResultHandler : TopkBlockResultHandler<C, use_sel> {
         std::vector<T> reservoir_dis;
         std::vector<TI> reservoir_ids;
 
-        explicit SingleResultHandler(ReservoirBlockResultHandler& hr)
-                : ReservoirTopN<C>(hr.k, hr.capacity, nullptr, nullptr),
-                  hr(hr) {}
+        explicit SingleResultHandler(ReservoirBlockResultHandler& hr_)
+                : ReservoirTopN<C>(hr_.k, hr_.capacity, nullptr, nullptr),
+                  hr(hr_) {}
 
         size_t qno;
 
@@ -533,11 +553,14 @@ struct ReservoirBlockResultHandler : TopkBlockResultHandler<C, use_sel> {
     }
 
     /// add results for query i0..i1 and j0..j1
-    void add_results(size_t j0, size_t j1, const T* dis_tab) {
+    void add_results(size_t j0, size_t j1, const T* dis_in) {
 #pragma omp parallel for
-        for (int64_t i = i0; i < i1; i++) {
-            ReservoirTopN<C>& reservoir = reservoirs[i - i0];
-            const T* dis_tab_i = dis_tab + (j1 - j0) * (i - i0) - j0;
+        for (int64_t i = static_cast<int64_t>(i0); i < static_cast<int64_t>(i1);
+             i++) {
+            ReservoirTopN<C>& reservoir =
+                    reservoirs[i - static_cast<int64_t>(i0)];
+            const T* dis_tab_i =
+                    dis_in + (j1 - j0) * (i - static_cast<int64_t>(i0)) - j0;
             for (size_t j = j0; j < j1; j++) {
                 T dis = dis_tab_i[j];
                 reservoir.add_result(dis, j);
@@ -547,9 +570,10 @@ struct ReservoirBlockResultHandler : TopkBlockResultHandler<C, use_sel> {
 
     /// series of results for queries i0..i1 is done
     void end_multiple() final {
-        // maybe parallel for
-        for (size_t i = i0; i < i1; i++) {
-            reservoirs[i - i0].to_result(
+#pragma omp parallel for schedule(static) if ((i1 - i0) * this->k >= 1024)
+        for (int64_t i = static_cast<int64_t>(i0); i < static_cast<int64_t>(i1);
+             i++) {
+            reservoirs[i - static_cast<int64_t>(i0)].to_result(
                     this->dis_tab + i * this->k, this->ids_tab + i * this->k);
         }
     }
@@ -567,14 +591,15 @@ struct RangeResultHandler : ResultHandlerT<C> {
 
     RangeQueryResult* qr = nullptr;
 
-    RangeResultHandler(RangeQueryResult* qr, T threshold) : qr(qr) {
-        this->threshold = threshold;
+    RangeResultHandler(RangeQueryResult* qr_, T threshold_) : qr(qr_) {
+        this->threshold = threshold_;
     }
 
     /// add one result for query i
     bool add_result(T dis, TI idx) final {
         if (C::cmp(threshold, dis)) {
             qr->add(dis, idx);
+            return true;
         }
         return false;
     }
@@ -591,12 +616,12 @@ struct RangeSearchBlockResultHandler : BlockResultHandler<C, use_sel> {
     T radius;
 
     RangeSearchBlockResultHandler(
-            RangeSearchResult* res,
-            float radius,
-            const IDSelector* sel = nullptr)
-            : BlockResultHandler<C, use_sel>(res->nq, sel),
-              res(res),
-              radius(radius) {}
+            RangeSearchResult* res_,
+            float radius_,
+            const IDSelector* sel_ = nullptr)
+            : BlockResultHandler<C, use_sel>(res_->nq, sel_),
+              res(res_),
+              radius(radius_) {}
 
     /******************************************************
      * API for 1 result at a time (each SingleResultHandler is
@@ -656,7 +681,7 @@ struct RangeSearchBlockResultHandler : BlockResultHandler<C, use_sel> {
         // it is a bit tricky to find the proper PartialResult structure
         // because the inner loop is on db not on queries.
 
-        if (pr < j0s.size() && j0 == j0s[pr]) {
+        if (static_cast<size_t>(pr) < j0s.size() && j0 == j0s[pr]) {
             pres = partial_results[pr];
             pr++;
         } else if (j0 == 0 && j0s.size() > 0) {
@@ -718,7 +743,7 @@ typename Consumer::T dispatch_knn_ResultHandler(
     if (k == 1) {                                                           \
         Top1BlockResultHandler<C, use_sel> res(nx, vals, ids, sel);         \
         return consumer.template f<>(res, args...);                         \
-    } else if (k < distance_compute_min_k_reservoir) {                      \
+    } else if (k < static_cast<size_t>(distance_compute_min_k_reservoir)) { \
         HeapBlockResultHandler<C, use_sel> res(nx, vals, ids, k, sel);      \
         return consumer.template f<>(res, args...);                         \
     } else {                                                                \