faiss 0.1.3 → 0.2.0

data/vendor/faiss/faiss/IndexBinary.h

@@ -120,7 +120,7 @@ struct IndexBinary {
    * @param x           input vectors to search, size n * d / 8
    * @param labels      output labels of the NNs, size n*k
    */
-  void assign(idx_t n, const uint8_t *x, idx_t *labels, idx_t k = 1);
+  void assign(idx_t n, const uint8_t *x, idx_t *labels, idx_t k = 1) const;
 
   /// Removes all elements from the database.
   virtual void reset() = 0;

data/vendor/faiss/faiss/IndexBinaryHash.cpp

@@ -18,16 +18,7 @@
 
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>
-
-#ifdef _MSC_VER
-#include <intrin.h>
-
-static inline int __builtin_ctzll(uint64_t x) {
-    unsigned long ret;
-    _BitScanForward64(&ret, x);
-    return (int)ret;
-}
-#endif // _MSC_VER
+#include <faiss/impl/platform_macros.h>
 
 namespace faiss {
 
@@ -145,8 +136,7 @@ struct KnnSearchResults {
 
     inline void add (float dis, idx_t id) {
         if (dis < heap_sim[0]) {
-            heap_pop<C> (k, heap_sim, heap_ids);
-            heap_push<C> (k, heap_sim, heap_ids, dis, id);
+            heap_replace_top<C> (k, heap_sim, heap_ids, dis, id);
         }
     }
 

data/vendor/faiss/faiss/IndexBinaryIVF.cpp

@@ -319,9 +319,8 @@ struct IVFBinaryScannerL2: BinaryInvertedListScanner {
         for (size_t j = 0; j < n; j++) {
             uint32_t dis = hc.hamming (codes);
             if (dis < simi[0]) {
-                heap_pop<C> (k, simi, idxi);
                 idx_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++;
             }
             codes += code_size;

data/vendor/faiss/faiss/IndexFlat.cpp

@@ -226,155 +226,7 @@ void IndexFlat::sa_decode (idx_t n, const uint8_t *bytes, float *x) const
 
 
 
-/***************************************************
- * IndexFlatL2BaseShift
- ***************************************************/
-
-IndexFlatL2BaseShift::IndexFlatL2BaseShift (idx_t d, size_t nshift, const float *shift):
-    IndexFlatL2 (d), shift (nshift)
-{
-    memcpy (this->shift.data(), shift, sizeof(float) * nshift);
-}
 
-void IndexFlatL2BaseShift::search (
-            idx_t n,
-            const float *x,
-            idx_t k,
-            float *distances,
-            idx_t *labels) const
-{
-    FAISS_THROW_IF_NOT (shift.size() == ntotal);
-
-    float_maxheap_array_t res = {
-        size_t(n), size_t(k), labels, distances};
-    knn_L2sqr_base_shift (x, xb.data(), d, n, ntotal, &res, shift.data());
-}
-
-
-
-/***************************************************
- * IndexRefineFlat
- ***************************************************/
-
-IndexRefineFlat::IndexRefineFlat (Index *base_index):
-    Index (base_index->d, base_index->metric_type),
-    refine_index (base_index->d, base_index->metric_type),
-    base_index (base_index), own_fields (false),
-    k_factor (1)
-{
-    is_trained = base_index->is_trained;
-    FAISS_THROW_IF_NOT_MSG (base_index->ntotal == 0,
-                      "base_index should be empty in the beginning");
-}
-
-IndexRefineFlat::IndexRefineFlat () {
-    base_index = nullptr;
-    own_fields = false;
-    k_factor = 1;
-}
-
-
-void IndexRefineFlat::train (idx_t n, const float *x)
-{
-    base_index->train (n, x);
-    is_trained = true;
-}
-
-void IndexRefineFlat::add (idx_t n, const float *x) {
-    FAISS_THROW_IF_NOT (is_trained);
-    base_index->add (n, x);
-    refine_index.add (n, x);
-    ntotal = refine_index.ntotal;
-}
-
-void IndexRefineFlat::reset ()
-{
-    base_index->reset ();
-    refine_index.reset ();
-    ntotal = 0;
-}
-
-namespace {
-typedef faiss::Index::idx_t idx_t;
-
-template<class C>
-static void reorder_2_heaps (
-      idx_t n,
-      idx_t k, idx_t *labels, float *distances,
-      idx_t k_base, const idx_t *base_labels, const float *base_distances)
-{
-#pragma omp parallel for
-    for (idx_t i = 0; i < n; i++) {
-        idx_t *idxo = labels + i * k;
-        float *diso = distances + i * k;
-        const idx_t *idxi = base_labels + i * k_base;
-        const float *disi = base_distances + i * k_base;
-
-        heap_heapify<C> (k, diso, idxo, disi, idxi, k);
-        if (k_base != k) { // add remaining elements
-            heap_addn<C> (k, diso, idxo, disi + k, idxi + k, k_base - k);
-        }
-        heap_reorder<C> (k, diso, idxo);
-    }
-}
-
-
-}
-
-
-void IndexRefineFlat::search (
-              idx_t n, const float *x, idx_t k,
-              float *distances, idx_t *labels) const
-{
-    FAISS_THROW_IF_NOT (is_trained);
-    idx_t k_base = idx_t (k * k_factor);
-    idx_t * base_labels = labels;
-    float * base_distances = distances;
-    ScopeDeleter<idx_t> del1;
-    ScopeDeleter<float> del2;
-
-
-    if (k != k_base) {
-        base_labels = new idx_t [n * k_base];
-        del1.set (base_labels);
-        base_distances = new float [n * k_base];
-        del2.set (base_distances);
-    }
-
-    base_index->search (n, x, k_base, base_distances, base_labels);
-
-    for (int i = 0; i < n * k_base; i++)
-        assert (base_labels[i] >= -1 &&
-                base_labels[i] < ntotal);
-
-    // compute refined distances
-    refine_index.compute_distance_subset (
-        n, x, k_base, base_distances, base_labels);
-
-    // sort and store result
-    if (metric_type == METRIC_L2) {
-        typedef CMax <float, idx_t> C;
-        reorder_2_heaps<C> (
-            n, k, labels, distances,
-            k_base, base_labels, base_distances);
-
-    } else if (metric_type == METRIC_INNER_PRODUCT) {
-        typedef CMin <float, idx_t> C;
-        reorder_2_heaps<C> (
-            n, k, labels, distances,
-            k_base, base_labels, base_distances);
-    } else {
-        FAISS_THROW_MSG("Metric type not supported");
-    }
-
-}
-
-
-
-IndexRefineFlat::~IndexRefineFlat ()
-{
-    if (own_fields) delete base_index;
-}
 
 /***************************************************
  * IndexFlat1D

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;