faiss 0.1.3 → 0.1.4

data/vendor/faiss/faiss/IndexFlat.h

@@ -93,57 +93,6 @@ struct IndexFlatL2:IndexFlat {
 };
 
 
-// same as an IndexFlatL2 but a value is subtracted from each distance
-struct IndexFlatL2BaseShift: IndexFlatL2 {
-    std::vector<float> shift;
-
-    IndexFlatL2BaseShift (idx_t d, size_t nshift, const float *shift);
-
-    void search(
-        idx_t n,
-        const float* x,
-        idx_t k,
-        float* distances,
-        idx_t* labels) const override;
-};
-
-
-/** Index that queries in a base_index (a fast one) and refines the
- *  results with an exact search, hopefully improving the results.
- */
-struct IndexRefineFlat: Index {
-
-    /// storage for full vectors
-    IndexFlat refine_index;
-
-    /// faster index to pre-select the vectors that should be filtered
-    Index *base_index;
-    bool own_fields;  ///< should the base index be deallocated?
-
-    /// factor between k requested in search and the k requested from
-    /// the base_index (should be >= 1)
-    float k_factor;
-
-    explicit IndexRefineFlat (Index *base_index);
-
-    IndexRefineFlat ();
-
-    void train(idx_t n, const float* x) override;
-
-    void add(idx_t n, const float* x) override;
-
-    void reset() override;
-
-    void search(
-        idx_t n,
-        const float* x,
-        idx_t k,
-        float* distances,
-        idx_t* labels) const override;
-
-    ~IndexRefineFlat() override;
-};
-
 
 /// optimized version for 1D "vectors".
 struct IndexFlat1D:IndexFlatL2 {

data/vendor/faiss/faiss/IndexHNSW.cpp

@@ -277,7 +277,7 @@ IndexHNSW::~IndexHNSW() {
 void IndexHNSW::train(idx_t n, const float* x)
 {
     FAISS_THROW_IF_NOT_MSG(storage,
-       "Please use IndexHSNWFlat (or variants) instead of IndexHNSW directly");
+       "Please use IndexHNSWFlat (or variants) instead of IndexHNSW directly");
     // hnsw structure does not require training
     storage->train (n, x);
     is_trained = true;
@@ -288,7 +288,7 @@ void IndexHNSW::search (idx_t n, const float *x, idx_t k,
 
 {
     FAISS_THROW_IF_NOT_MSG(storage,
-       "Please use IndexHSNWFlat (or variants) instead of IndexHNSW directly");
+       "Please use IndexHNSWFlat (or variants) instead of IndexHNSW directly");
     size_t n1 = 0, n2 = 0, n3 = 0, ndis = 0, nreorder = 0;
 
     idx_t check_period = InterruptCallback::get_period_hint (
@@ -352,7 +352,7 @@ void IndexHNSW::search (idx_t n, const float *x, idx_t k,
 void IndexHNSW::add(idx_t n, const float *x)
 {
     FAISS_THROW_IF_NOT_MSG(storage,
-       "Please use IndexHSNWFlat (or variants) instead of IndexHNSW directly");
+       "Please use IndexHNSWFlat (or variants) instead of IndexHNSW directly");
     FAISS_THROW_IF_NOT(is_trained);
     int n0 = ntotal;
     storage->add(n, x);
@@ -1003,8 +1003,7 @@ int search_from_candidates_2(const HNSW & hnsw,
                     if (nres < k) {
                         faiss::maxheap_push (++nres, D, I, d, v1);
                     } else if (d < D[0]) {
-                        faiss::maxheap_pop (nres--, D, I);
-                        faiss::maxheap_push (++nres, D, I, d, v1);
+                        faiss::maxheap_replace_top (nres, D, I, d, v1);
                     }
                 }
                 vt.visited[v1] = vt.visno + 1;

data/vendor/faiss/faiss/IndexIVF.cpp

@@ -88,12 +88,19 @@ void Level1Quantizer::train_q1 (size_t n, const float *x, bool verbose, MetricTy
         }
         quantizer->is_trained = true;
     } else if (quantizer_trains_alone == 2) {
-        if (verbose)
+        if (verbose) {
             printf (
                 "Training L2 quantizer on %zd vectors in %zdD%s\n",
                 n, d,
                 clustering_index ? "(user provided index)" : "");
-        FAISS_THROW_IF_NOT (metric_type == METRIC_L2);
+        }
+        // also accept spherical centroids because in that case
+        // L2 and IP are equivalent
+        FAISS_THROW_IF_NOT (
+            metric_type == METRIC_L2 ||
+            (metric_type == METRIC_INNER_PRODUCT && cp.spherical)
+        );
+
         Clustering clus (d, nlist, cp);
         if (!clustering_index) {
             IndexFlatL2 assigner (d);
@@ -263,23 +270,76 @@ void IndexIVF::set_direct_map_type (DirectMap::Type type)
     direct_map.set_type (type, invlists, ntotal);
 }
 
-
+/** It is a sad fact of software that a conceptually simple function like this
+ * becomes very complex when you factor in several ways of parallelizing +
+ * interrupt/error handling + collecting stats + min/max collection. The
+ * codepath that is used 95% of time is the one for parallel_mode = 0 */
 void IndexIVF::search (idx_t n, const float *x, idx_t k,
                          float *distances, idx_t *labels) const
 {
-    std::unique_ptr<idx_t[]> idx(new idx_t[n * nprobe]);
-    std::unique_ptr<float[]> coarse_dis(new float[n * nprobe]);
 
-    double t0 = getmillisecs();
-    quantizer->search (n, x, nprobe, coarse_dis.get(), idx.get());
-    indexIVF_stats.quantization_time += getmillisecs() - t0;
 
-    t0 = getmillisecs();
-    invlists->prefetch_lists (idx.get(), n * nprobe);
+    // search function for a subset of queries
+    auto sub_search_func = [this, k]
+            (idx_t n, const float *x, float *distances, idx_t *labels,
+             IndexIVFStats *ivf_stats) {
+
+        std::unique_ptr<idx_t[]> idx(new idx_t[n * nprobe]);
+        std::unique_ptr<float[]> coarse_dis(new float[n * nprobe]);
+
+        double t0 = getmillisecs();
+        quantizer->search (n, x, nprobe, coarse_dis.get(), idx.get());
+
+        double t1 = getmillisecs();
+        invlists->prefetch_lists (idx.get(), n * nprobe);
+
+        search_preassigned (n, x, k, idx.get(), coarse_dis.get(),
+                            distances, labels, false, nullptr, ivf_stats);
+        double t2 = getmillisecs();
+        ivf_stats->quantization_time += t1 - t0;
+        ivf_stats->search_time += t2 - t0;
+    };
+
+
+    if ((parallel_mode & ~PARALLEL_MODE_NO_HEAP_INIT) == 0) {
+        int nt = std::min(omp_get_max_threads(), int(n));
+        std::vector<IndexIVFStats> stats(nt);
+        std::mutex exception_mutex;
+        std::string exception_string;
+
+#pragma omp parallel for if (nt > 1)
+        for(idx_t slice = 0; slice < nt; slice++) {
+            IndexIVFStats local_stats;
+            idx_t i0 = n * slice / nt;
+            idx_t i1 = n * (slice + 1) / nt;
+            if (i1 > i0) {
+                try {
+                    sub_search_func(
+                            i1 - i0, x + i0 * d,
+                            distances + i0 * k, labels + i0 * k,
+                            &stats[slice]
+                    );
+                } catch(const std::exception & e) {
+                    std::lock_guard<std::mutex> lock(exception_mutex);
+                    exception_string = e.what();
+                }
+            }
+        }
+
+        if (!exception_string.empty()) {
+            FAISS_THROW_MSG (exception_string.c_str());
+        }
+
+        // collect stats
+        for(idx_t slice = 0; slice < nt; slice++) {
+            indexIVF_stats.add(stats[slice]);
+        }
+    } else {
+        // handle paralellization at level below (or don't run in parallel at all)
+        sub_search_func(n, x, distances, labels, &indexIVF_stats);
+    }
+
 
-    search_preassigned (n, x, k, idx.get(), coarse_dis.get(),
-                        distances, labels, false);
-    indexIVF_stats.search_time += getmillisecs() - t0;
 }
 
 
@@ -288,7 +348,8 @@ void IndexIVF::search_preassigned (idx_t n, const float *x, idx_t k,
                                    const float *coarse_dis ,
                                    float *distances, idx_t *labels,
                                    bool store_pairs,
-                                   const IVFSearchParameters *params) const
+                                   const IVFSearchParameters *params,
+                                   IndexIVFStats *ivf_stats) const
 {
     long nprobe = params ? params->nprobe : this->nprobe;
     long max_codes = params ? params->max_codes : this->max_codes;
@@ -305,13 +366,12 @@ void IndexIVF::search_preassigned (idx_t n, const float *x, idx_t k,
     int pmode = this->parallel_mode & ~PARALLEL_MODE_NO_HEAP_INIT;
     bool do_heap_init = !(this->parallel_mode & PARALLEL_MODE_NO_HEAP_INIT);
 
-    // don't start parallel section if single query
     bool do_parallel = omp_get_max_threads() >= 2 && (
-            pmode == 0 ? n > 1 :
+            pmode == 0 ? false :
+            pmode == 3 ? n > 1 :
             pmode == 1 ? nprobe > 1 :
             nprobe * n > 1);
 
-
 #pragma omp parallel if(do_parallel) reduction(+: nlistv, ndis, nheap)
     {
         InvertedListScanner *scanner = get_InvertedListScanner(store_pairs);
@@ -409,7 +469,7 @@ void IndexIVF::search_preassigned (idx_t n, const float *x, idx_t k,
          * Actual loops, depending on parallel_mode
          ****************************************************/
 
-        if (pmode == 0) {
+        if (pmode == 0 || pmode == 3) {
 
 #pragma omp for
             for (idx_t i = 0; i < n; i++) {
@@ -527,11 +587,12 @@ void IndexIVF::search_preassigned (idx_t n, const float *x, idx_t k,
         }
     }
 
-    indexIVF_stats.nq += n;
-    indexIVF_stats.nlist += nlistv;
-    indexIVF_stats.ndis += ndis;
-    indexIVF_stats.nheap_updates += nheap;
-
+    if (ivf_stats) {
+        ivf_stats->nq += n;
+        ivf_stats->nlist += nlistv;
+        ivf_stats->ndis += ndis;
+        ivf_stats->nheap_updates += nheap;
+    }
 }
 
 
@@ -551,7 +612,7 @@ void IndexIVF::range_search (idx_t nx, const float *x, float radius,
     invlists->prefetch_lists (keys.get(), nx * nprobe);
 
     range_search_preassigned (nx, x, radius, keys.get (), coarse_dis.get (),
-                              result);
+                              result, false, nullptr, &indexIVF_stats);
 
     indexIVF_stats.search_time += getmillisecs() - t0;
 }
@@ -561,7 +622,8 @@ void IndexIVF::range_search_preassigned (
          const idx_t *keys, const float *coarse_dis,
          RangeSearchResult *result,
          bool store_pairs,
-         const IVFSearchParameters *params) const
+         const IVFSearchParameters *params,
+         IndexIVFStats *stats) const
 {
     long nprobe = params ? params->nprobe : this->nprobe;
     long max_codes = params ? params->max_codes : this->max_codes;
@@ -574,7 +636,15 @@ void IndexIVF::range_search_preassigned (
 
     std::vector<RangeSearchPartialResult *> all_pres (omp_get_max_threads());
 
-#pragma omp parallel reduction(+: nlistv, ndis)
+    int pmode = this->parallel_mode & ~PARALLEL_MODE_NO_HEAP_INIT;
+    // don't start parallel section if single query
+    bool do_parallel = omp_get_max_threads() >= 2 && (
+            pmode == 3 ? false :
+            pmode == 0 ? nx > 1 :
+            pmode == 1 ? nprobe > 1 :
+            nprobe * nx > 1);
+
+#pragma omp parallel if(do_parallel) reduction(+: nlistv, ndis)
     {
         RangeSearchPartialResult pres(result);
         std::unique_ptr<InvertedListScanner> scanner
@@ -680,9 +750,11 @@ void IndexIVF::range_search_preassigned (
         }
     }
 
-    indexIVF_stats.nq += nx;
-    indexIVF_stats.nlist += nlistv;
-    indexIVF_stats.ndis += ndis;
+    if (stats) {
+        stats->nq += nx;
+        stats->nlist += nlistv;
+        stats->ndis += ndis;
+    }
 }
 
 
@@ -879,11 +951,15 @@ void IndexIVF::replace_invlists (InvertedLists *il, bool own)
 {
     if (own_invlists) {
         delete invlists;
+        invlists = nullptr;
     }
     // FAISS_THROW_IF_NOT (ntotal == 0);
     if (il) {
-        FAISS_THROW_IF_NOT (il->nlist == nlist &&
-                            il->code_size == code_size);
+        FAISS_THROW_IF_NOT (il->nlist == nlist);
+        FAISS_THROW_IF_NOT (
+            il->code_size == code_size ||
+            il->code_size == InvertedLists::INVALID_CODE_SIZE
+        );
     }
     invlists = il;
     own_invlists = own;
@@ -971,6 +1047,17 @@ void IndexIVFStats::reset()
     memset ((void*)this, 0, sizeof (*this));
 }
 
+void IndexIVFStats::add (const IndexIVFStats & other)
+{
+    nq += other.nq;
+    nlist += other.nlist;
+    ndis += other.ndis;
+    nheap_updates += other.nheap_updates;
+    quantization_time += other.quantization_time;
+    search_time += other.search_time;
+
+}
+
 
 IndexIVFStats indexIVF_stats;
 

data/vendor/faiss/faiss/IndexIVF.h

@@ -16,8 +16,8 @@
 #include <stdint.h>
 
 #include <faiss/Index.h>
-#include <faiss/InvertedLists.h>
-#include <faiss/DirectMap.h>
+#include <faiss/invlists/InvertedLists.h>
+#include <faiss/invlists/DirectMap.h>
 #include <faiss/Clustering.h>
 #include <faiss/impl/platform_macros.h>
 #include <faiss/utils/Heap.h>
@@ -76,6 +76,7 @@ struct IVFSearchParameters {
 
 
 struct InvertedListScanner;
+struct IndexIVFStats;
 
 /** Index based on a inverted file (IVF)
  *
@@ -109,9 +110,10 @@ struct IndexIVF: Index, Level1Quantizer {
 
     /** Parallel mode determines how queries are parallelized with OpenMP
      *
-     * 0 (default): parallelize over queries
+     * 0 (default): split over queries
      * 1: parallelize over inverted lists
      * 2: parallelize over both
+     * 3: split over queries with a finer granularity
      *
      * PARALLEL_MODE_NO_HEAP_INIT: binary or with the previous to
      * prevent the heap to be initialized and finalized
@@ -178,14 +180,16 @@ struct IndexIVF: Index, Level1Quantizer {
      *                     instead in upper/lower 32 bit of result,
      *                     instead of ids (used for reranking).
      * @param params used to override the object's search parameters
+     * @param stats  search stats to be updated (can be null)
      */
-    virtual void search_preassigned (idx_t n, const float *x, idx_t k,
-                                     const idx_t *assign,
-                                     const float *centroid_dis,
-                                     float *distances, idx_t *labels,
-                                     bool store_pairs,
-                                     const IVFSearchParameters *params=nullptr
-                                     ) const;
+    virtual void search_preassigned (
+            idx_t n, const float *x, idx_t k,
+            const idx_t *assign, const float *centroid_dis,
+            float *distances, idx_t *labels,
+            bool store_pairs,
+            const IVFSearchParameters *params=nullptr,
+            IndexIVFStats *stats=nullptr
+            ) const;
 
     /** assign the vectors, then call search_preassign */
     void search (idx_t n, const float *x, idx_t k,
@@ -194,11 +198,13 @@ struct IndexIVF: Index, Level1Quantizer {
     void range_search (idx_t n, const float* x, float radius,
                        RangeSearchResult* result) const override;
 
-    void range_search_preassigned(idx_t nx, const float *x, float radius,
-                                  const idx_t *keys, const float *coarse_dis,
-                                  RangeSearchResult *result,
-                                  bool store_pairs=false,
-                                  const IVFSearchParameters *params=nullptr) const;
+    void range_search_preassigned(
+            idx_t nx, const float *x, float radius,
+            const idx_t *keys, const float *coarse_dis,
+            RangeSearchResult *result,
+            bool store_pairs=false,
+            const IVFSearchParameters *params=nullptr,
+            IndexIVFStats *stats=nullptr) const;
 
     /// get a scanner for this index (store_pairs means ignore labels)
     virtual InvertedListScanner *get_InvertedListScanner (
@@ -364,6 +370,7 @@ struct IndexIVFStats {
 
     IndexIVFStats () {reset (); }
     void reset ();
+    void add (const IndexIVFStats & other);
 };
 
 // global var that collects them all

data/vendor/faiss/faiss/IndexIVFFlat.cpp

@@ -159,9 +159,8 @@ struct IVFFlatScanner: InvertedListScanner {
             float dis = metric == METRIC_INNER_PRODUCT ?
                 fvec_inner_product (xi, yj, d) : fvec_L2sqr (xi, yj, d);
             if (C::cmp (simi[0], dis)) {
-                heap_pop<C> (k, simi, idxi);
                 int64_t id = store_pairs ? lo_build (list_no, j) : ids[j];
-                heap_push<C> (k, simi, idxi, dis, id);
+                heap_replace_top<C> (k, simi, idxi, dis, id);
                 nup++;
             }
         }
@@ -317,7 +316,8 @@ void IndexIVFFlatDedup::search_preassigned (
            const float *centroid_dis,
            float *distances, idx_t *labels,
            bool store_pairs,
-           const IVFSearchParameters *params) const
+           const IVFSearchParameters *params,
+           IndexIVFStats *stats) const
 {
     FAISS_THROW_IF_NOT_MSG (
            !store_pairs, "store_pairs not supported in IVFDedup");

data/vendor/faiss/faiss/IndexIVFFlat.h

@@ -77,7 +77,8 @@ struct IndexIVFFlatDedup: IndexIVFFlat {
                              const float *centroid_dis,
                              float *distances, idx_t *labels,
                              bool store_pairs,
-                             const IVFSearchParameters *params=nullptr
+                             const IVFSearchParameters *params=nullptr,
+                             IndexIVFStats *stats=nullptr
                              ) const override;
 
     size_t remove_ids(const IDSelector& sel) override;