faiss 0.1.1 → 0.1.2

data/vendor/faiss/IVFlib.cpp

@@ -13,6 +13,7 @@
 
 #include <faiss/IndexPreTransform.h>
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/MetaIndexes.h>
 
 
 
@@ -56,17 +57,35 @@ void check_compatible_for_merge (const Index * index0,
 
 }
 
-const IndexIVF * extract_index_ivf (const Index * index)
+const IndexIVF * try_extract_index_ivf (const Index * index)
 {
     if (auto *pt =
         dynamic_cast<const IndexPreTransform *>(index)) {
         index = pt->index;
     }
 
+    if (auto *idmap =
+        dynamic_cast<const IndexIDMap *>(index)) {
+        index = idmap->index;
+    }
+    if (auto *idmap =
+        dynamic_cast<const IndexIDMap2 *>(index)) {
+        index = idmap->index;
+    }
+
     auto *ivf = dynamic_cast<const IndexIVF *>(index);
 
-    FAISS_THROW_IF_NOT (ivf);
+    return ivf;
+}
 
+IndexIVF * try_extract_index_ivf (Index * index) {
+    return const_cast<IndexIVF*> (try_extract_index_ivf ((const Index*)(index)));
+}
+
+const IndexIVF * extract_index_ivf (const Index * index)
+{
+    const IndexIVF *ivf = try_extract_index_ivf (index);
+    FAISS_THROW_IF_NOT (ivf);
     return ivf;
 }
 
@@ -74,6 +93,7 @@ IndexIVF * extract_index_ivf (Index * index) {
     return const_cast<IndexIVF*> (extract_index_ivf ((const Index*)(index)));
 }
 
+
 void merge_into(faiss::Index *index0, faiss::Index *index1, bool shift_ids) {
 
     check_compatible_for_merge (index0, index1);
@@ -146,8 +166,8 @@ void search_and_return_centroids(faiss::Index *index,
             if (result_centroid_ids)
                 result_centroid_ids[i] = -1;
         } else {
-            long list_no = label >> 32;
-            long list_index = label & 0xffffffff;
+            long list_no = lo_listno (label);
+            long list_index = lo_offset (label);
             if (result_centroid_ids)
                 result_centroid_ids[i] = list_no;
             labels[i] = index_ivf->invlists->get_single_id(list_no, list_index);

data/vendor/faiss/IVFlib.h

@@ -35,6 +35,10 @@ void check_compatible_for_merge (const Index * index1,
 const IndexIVF * extract_index_ivf (const Index * index);
 IndexIVF * extract_index_ivf (Index * index);
 
+/// same as above but returns nullptr instead of throwing on failure
+const IndexIVF * try_extract_index_ivf (const Index * index);
+IndexIVF * try_extract_index_ivf (Index * index);
+
 /** Merge index1 into index0. Works on IndexIVF's and IndexIVF's
  *  embedded in a IndexPreTransform. On output, the index1 is empty.
  *

data/vendor/faiss/Index.h

@@ -10,7 +10,7 @@
 #ifndef FAISS_INDEX_H
 #define FAISS_INDEX_H
 
-
+#include <faiss/MetricType.h>
 #include <cstdio>
 #include <typeinfo>
 #include <string>
@@ -18,7 +18,7 @@
 
 #define FAISS_VERSION_MAJOR 1
 #define FAISS_VERSION_MINOR 6
-#define FAISS_VERSION_PATCH 1
+#define FAISS_VERSION_PATCH 3
 
 /**
  * @namespace faiss
@@ -39,34 +39,15 @@
 
 namespace faiss {
 
-
-/// Some algorithms support both an inner product version and a L2 search version.
-enum MetricType {
-    METRIC_INNER_PRODUCT = 0,  ///< maximum inner product search
-    METRIC_L2 = 1,             ///< squared L2 search
-    METRIC_L1,                 ///< L1 (aka cityblock)
-    METRIC_Linf,               ///< infinity distance
-    METRIC_Lp,                 ///< L_p distance, p is given by metric_arg
-
-    /// some additional metrics defined in scipy.spatial.distance
-    METRIC_Canberra = 20,
-    METRIC_BrayCurtis,
-    METRIC_JensenShannon,
-
-};
-
-
 /// Forward declarations see AuxIndexStructures.h
 struct IDSelector;
 struct RangeSearchResult;
 struct DistanceComputer;
 
-/** Abstract structure for an index
- *
- * Supports adding vertices and searching them.
+/** Abstract structure for an index, supports adding vectors and searching them.
  *
- * Currently only asymmetric queries are supported:
- * database-to-database queries are not implemented.
+ * All vectors provided at add or search time are 32-bit float arrays,
+ * although the internal representation may vary.
  */
 struct Index {
     using idx_t = int64_t;  ///< all indices are this type

data/vendor/faiss/Index2Layer.cpp

@@ -42,7 +42,6 @@
 
 namespace faiss {
 
-using idx_t = Index::idx_t;
 
 /*************************************
  * Index2Layer implementation

data/vendor/faiss/IndexBinary.h

@@ -99,9 +99,13 @@ struct IndexBinary {
 
   /** Query n vectors of dimension d to the index.
    *
-   * return all vectors with distance < radius. Note that many
-   * indexes do not implement the range_search (only the k-NN search
-   * is mandatory).
+   * return all vectors with distance < radius. Note that many indexes
+   * do not implement the range_search (only the k-NN search is
+   * mandatory). The distances are converted to float to reuse the
+   * RangeSearchResult structure, but they are integer. By convention,
+   * only distances < radius (strict comparison) are returned,
+   * ie. radius = 0 does not return any result and 1 returns only
+   * exact same vectors.
    *
    * @param x           input vectors to search, size n * d / 8
    * @param radius      search radius

data/vendor/faiss/IndexBinaryFlat.cpp

@@ -79,5 +79,10 @@ void IndexBinaryFlat::reconstruct(idx_t key, uint8_t *recons) const {
   memcpy(recons, &(xb[code_size * key]), sizeof(*recons) * code_size);
 }
 
+void IndexBinaryFlat::range_search(idx_t n, const uint8_t *x, int radius,
+                   RangeSearchResult *result) const
+{
+    hamming_range_search (x, xb.data(), n, ntotal, radius, code_size, result);
+}
 
 }  // namespace faiss

data/vendor/faiss/IndexBinaryFlat.h

@@ -38,6 +38,9 @@ struct IndexBinaryFlat : IndexBinary {
   void search(idx_t n, const uint8_t *x, idx_t k,
               int32_t *distances, idx_t *labels) const override;
 
+  void range_search(idx_t n, const uint8_t *x, int radius,
+                   RangeSearchResult *result) const override;
+
   void reconstruct(idx_t key, uint8_t *recons) const override;
 
   /** Remove some ids. Note that because of the indexing structure,

data/vendor/faiss/IndexBinaryHash.cpp

@@ -0,0 +1,492 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// Copyright 2004-present Facebook. All Rights Reserved
+// -*- c++ -*-
+
+#include <faiss/IndexBinaryHash.h>
+
+#include <cstdio>
+#include <memory>
+
+#include <faiss/utils/hamming.h>
+#include <faiss/utils/utils.h>
+
+#include <faiss/impl/AuxIndexStructures.h>
+#include <faiss/impl/FaissAssert.h>
+
+
+namespace faiss {
+
+void IndexBinaryHash::InvertedList::add (
+        idx_t id, size_t code_size, const uint8_t *code)
+{
+    ids.push_back(id);
+    vecs.insert(vecs.end(), code, code + code_size);
+}
+
+IndexBinaryHash::IndexBinaryHash(int d, int b):
+    IndexBinary(d), b(b), nflip(0)
+{
+    is_trained = true;
+}
+
+IndexBinaryHash::IndexBinaryHash(): b(0), nflip(0)
+{
+    is_trained = true;
+}
+
+void IndexBinaryHash::reset()
+{
+    invlists.clear();
+    ntotal = 0;
+}
+
+
+void IndexBinaryHash::add(idx_t n, const uint8_t *x)
+{
+    add_with_ids(n, x, nullptr);
+}
+
+void IndexBinaryHash::add_with_ids(idx_t n, const uint8_t *x, const idx_t *xids)
+{
+    uint64_t mask = ((uint64_t)1 << b) - 1;
+    // simplistic add function. Cannot really be parallelized.
+
+    for (idx_t i = 0; i < n; i++) {
+        idx_t id = xids ? xids[i] : ntotal + i;
+        const uint8_t * xi = x + i * code_size;
+        idx_t hash = *((uint64_t*)xi) & mask;
+        invlists[hash].add(id, code_size, xi);
+    }
+    ntotal += n;
+}
+
+namespace {
+
+
+/** Enumerate all bit vectors of size nbit with up to maxflip 1s
+ * test in P127257851 P127258235
+ */
+struct FlipEnumerator {
+    int nbit, nflip, maxflip;
+    uint64_t mask, x;
+
+    FlipEnumerator (int nbit, int maxflip): nbit(nbit), maxflip(maxflip) {
+        nflip = 0;
+        mask = 0;
+        x = 0;
+    }
+
+    bool next() {
+        if (x == mask) {
+            if (nflip == maxflip) {
+                return false;
+            }
+            // increase Hamming radius
+            nflip++;
+            mask = (((uint64_t)1 << nflip) - 1);
+            x = mask << (nbit - nflip);
+            return true;
+        }
+
+        int i = __builtin_ctzll(x);
+
+        if (i > 0) {
+            x ^= (uint64_t)3 << (i - 1);
+        } else {
+            // nb of LSB 1s
+            int n1 = __builtin_ctzll(~x);
+            // clear them
+            x &= ((uint64_t)(-1) << n1);
+            int n2 = __builtin_ctzll(x);
+            x ^= (((uint64_t)1 << (n1 + 2)) - 1) << (n2 - n1 - 1);
+        }
+        return true;
+    }
+
+};
+
+using idx_t = Index::idx_t;
+
+
+struct RangeSearchResults {
+    int radius;
+    RangeQueryResult &qres;
+
+    inline void add (float dis, idx_t id) {
+        if (dis < radius) {
+            qres.add (dis, id);
+        }
+    }
+
+};
+
+struct KnnSearchResults {
+    // heap params
+    idx_t k;
+    int32_t * heap_sim;
+    idx_t * heap_ids;
+
+    using C = CMax<int, idx_t>;
+
+    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);
+        }
+    }
+
+};
+
+template<class HammingComputer, class SearchResults>
+void
+search_single_query_template(const IndexBinaryHash & index, const uint8_t *q,
+                    SearchResults &res,
+                    size_t &n0, size_t &nlist, size_t &ndis)
+{
+    size_t code_size = index.code_size;
+    uint64_t mask = ((uint64_t)1 << index.b) - 1;
+    uint64_t qhash = *((uint64_t*)q) & mask;
+    HammingComputer hc (q, code_size);
+    FlipEnumerator fe(index.b, index.nflip);
+
+    // loop over neighbors that are at most at nflip bits
+    do {
+        uint64_t hash = qhash ^ fe.x;
+        auto it = index.invlists.find (hash);
+
+        if (it == index.invlists.end()) {
+            continue;
+        }
+
+        const IndexBinaryHash::InvertedList &il = it->second;
+
+        size_t nv = il.ids.size();
+
+        if (nv == 0) {
+            n0++;
+        } else {
+            const uint8_t *codes = il.vecs.data();
+            for (size_t i = 0; i < nv; i++) {
+                int dis = hc.hamming (codes);
+                res.add(dis, il.ids[i]);
+                codes += code_size;
+            }
+            ndis += nv;
+            nlist++;
+        }
+    } while(fe.next());
+}
+
+template<class SearchResults>
+void
+search_single_query(const IndexBinaryHash & index, const uint8_t *q,
+                    SearchResults &res,
+                    size_t &n0, size_t &nlist, size_t &ndis)
+{
+#define HC(name) search_single_query_template<name>(index, q, res, n0, nlist, ndis);
+    switch(index.code_size) {
+    case 4: HC(HammingComputer4); break;
+    case 8: HC(HammingComputer8); break;
+    case 16: HC(HammingComputer16); break;
+    case 20: HC(HammingComputer20); break;
+    case 32: HC(HammingComputer32); break;
+    default:
+        if (index.code_size % 8 == 0) {
+            HC(HammingComputerM8);
+        } else {
+            HC(HammingComputerDefault);
+        }
+    }
+#undef HC
+}
+
+
+} // anonymous namespace
+
+
+
+void IndexBinaryHash::range_search(idx_t n, const uint8_t *x, int radius,
+                                   RangeSearchResult *result) const
+{
+
+    size_t nlist = 0, ndis = 0, n0 = 0;
+
+#pragma omp parallel if(n > 100) reduction(+: ndis, n0, nlist)
+    {
+        RangeSearchPartialResult pres (result);
+
+#pragma omp for
+        for (size_t i = 0; i < n; i++) { // loop queries
+            RangeQueryResult & qres = pres.new_result (i);
+            RangeSearchResults res = {radius, qres};
+            const uint8_t *q = x + i * code_size;
+
+            search_single_query (*this, q, res, n0, nlist, ndis);
+
+        }
+        pres.finalize ();
+    }
+    indexBinaryHash_stats.nq += n;
+    indexBinaryHash_stats.n0 += n0;
+    indexBinaryHash_stats.nlist += nlist;
+    indexBinaryHash_stats.ndis += ndis;
+}
+
+void IndexBinaryHash::search(idx_t n, const uint8_t *x, idx_t k,
+                             int32_t *distances, idx_t *labels) const
+{
+
+    using HeapForL2 = CMax<int32_t, idx_t>;
+    size_t nlist = 0, ndis = 0, n0 = 0;
+
+#pragma omp parallel for if(n > 100) reduction(+: nlist, ndis, n0)
+    for (size_t i = 0; i < n; i++) {
+        int32_t * simi = distances + k * i;
+        idx_t * idxi = labels + k * i;
+
+        heap_heapify<HeapForL2> (k, simi, idxi);
+        KnnSearchResults res = {k, simi, idxi};
+        const uint8_t *q = x + i * code_size;
+
+        search_single_query (*this, q, res, n0, nlist, ndis);
+
+    }
+    indexBinaryHash_stats.nq += n;
+    indexBinaryHash_stats.n0 += n0;
+    indexBinaryHash_stats.nlist += nlist;
+    indexBinaryHash_stats.ndis += ndis;
+}
+
+size_t IndexBinaryHash::hashtable_size() const
+{
+    return invlists.size();
+}
+
+
+void IndexBinaryHash::display() const
+{
+    for (auto it = invlists.begin(); it != invlists.end(); ++it) {
+        printf("%ld: [", it->first);
+        const std::vector<idx_t> & v = it->second.ids;
+        for (auto x: v) {
+            printf("%ld ", 0 + x);
+        }
+        printf("]\n");
+
+    }
+}
+
+
+void IndexBinaryHashStats::reset()
+{
+    memset ((void*)this, 0, sizeof (*this));
+}
+
+IndexBinaryHashStats indexBinaryHash_stats;
+
+/*******************************************************
+ * IndexBinaryMultiHash implementation
+ ******************************************************/
+
+
+IndexBinaryMultiHash::IndexBinaryMultiHash(int d, int nhash, int b):
+    IndexBinary(d),
+    storage(new IndexBinaryFlat(d)), own_fields(true),
+    maps(nhash), nhash(nhash), b(b), nflip(0)
+{
+    FAISS_THROW_IF_NOT(nhash * b <= d);
+}
+
+IndexBinaryMultiHash::IndexBinaryMultiHash():
+    storage(nullptr), own_fields(true),
+    nhash(0), b(0), nflip(0)
+{}
+
+IndexBinaryMultiHash::~IndexBinaryMultiHash()
+{
+    if (own_fields) {
+        delete storage;
+    }
+}
+
+
+void IndexBinaryMultiHash::reset()
+{
+    storage->reset();
+    ntotal = 0;
+    for(auto map: maps) {
+        map.clear();
+    }
+}
+
+void IndexBinaryMultiHash::add(idx_t n, const uint8_t *x)
+{
+    storage->add(n, x);
+    // populate maps
+    uint64_t mask = ((uint64_t)1 << b) - 1;
+
+    for(idx_t i = 0; i < n; i++) {
+        const uint8_t *xi = x + i * code_size;
+        int ho = 0;
+        for(int h = 0; h < nhash; h++) {
+            uint64_t hash = *(uint64_t*)(xi + (ho >> 3)) >> (ho & 7);
+            hash &= mask;
+            maps[h][hash].push_back(i + ntotal);
+            ho += b;
+        }
+    }
+    ntotal += n;
+}
+
+
+namespace {
+
+template <class HammingComputer, class SearchResults>
+static
+void verify_shortlist(
+        const IndexBinaryFlat & index,
+        const uint8_t * q,
+        const std::unordered_set<Index::idx_t> & shortlist,
+        SearchResults &res)
+{
+    size_t code_size = index.code_size;
+    size_t nlist = 0, ndis = 0, n0 = 0;
+
+    HammingComputer hc (q, code_size);
+    const uint8_t *codes = index.xb.data();
+
+    for (auto i: shortlist) {
+        int dis = hc.hamming (codes + i * code_size);
+        res.add(dis, i);
+    }
+}
+
+template<class SearchResults>
+void
+search_1_query_multihash(const IndexBinaryMultiHash & index, const uint8_t *xi,
+                         SearchResults &res,
+                         size_t &n0, size_t &nlist, size_t &ndis)
+{
+
+    std::unordered_set<idx_t> shortlist;
+    int b = index.b;
+    uint64_t mask = ((uint64_t)1 << b) - 1;
+
+    int ho = 0;
+    for(int h = 0; h < index.nhash; h++) {
+        uint64_t qhash = *(uint64_t*)(xi + (ho >> 3)) >> (ho & 7);
+        qhash &= mask;
+        const IndexBinaryMultiHash::Map & map = index.maps[h];
+
+        FlipEnumerator fe(index.b, index.nflip);
+        // loop over neighbors that are at most at nflip bits
+        do {
+            uint64_t hash = qhash ^ fe.x;
+            auto it = map.find (hash);
+
+            if (it != map.end()) {
+                const std::vector<idx_t> & v = it->second;
+                for (auto i: v) {
+                    shortlist.insert(i);
+                }
+                nlist++;
+            } else {
+                n0++;
+            }
+        } while(fe.next());
+
+        ho += b;
+    }
+    ndis += shortlist.size();
+
+    // verify shortlist
+
+#define HC(name) verify_shortlist<name> (*index.storage, xi, shortlist, res)
+    switch(index.code_size) {
+    case 4: HC(HammingComputer4); break;
+    case 8: HC(HammingComputer8); break;
+    case 16: HC(HammingComputer16); break;
+    case 20: HC(HammingComputer20); break;
+    case 32: HC(HammingComputer32); break;
+    default:
+        if (index.code_size % 8 == 0) {
+            HC(HammingComputerM8);
+        } else {
+            HC(HammingComputerDefault);
+        }
+    }
+#undef HC
+}
+
+} // anonymous namespace
+
+void IndexBinaryMultiHash::range_search(idx_t n, const uint8_t *x, int radius,
+                                   RangeSearchResult *result) const
+{
+
+    size_t nlist = 0, ndis = 0, n0 = 0;
+
+#pragma omp parallel if(n > 100) reduction(+: ndis, n0, nlist)
+    {
+        RangeSearchPartialResult pres (result);
+
+#pragma omp for
+        for (size_t i = 0; i < n; i++) { // loop queries
+            RangeQueryResult & qres = pres.new_result (i);
+            RangeSearchResults res = {radius, qres};
+            const uint8_t *q = x + i * code_size;
+
+            search_1_query_multihash (*this, q, res, n0, nlist, ndis);
+
+        }
+        pres.finalize ();
+    }
+    indexBinaryHash_stats.nq += n;
+    indexBinaryHash_stats.n0 += n0;
+    indexBinaryHash_stats.nlist += nlist;
+    indexBinaryHash_stats.ndis += ndis;
+}
+
+void IndexBinaryMultiHash::search(idx_t n, const uint8_t *x, idx_t k,
+                             int32_t *distances, idx_t *labels) const
+{
+
+    using HeapForL2 = CMax<int32_t, idx_t>;
+    size_t nlist = 0, ndis = 0, n0 = 0;
+
+#pragma omp parallel for if(n > 100) reduction(+: nlist, ndis, n0)
+    for (size_t i = 0; i < n; i++) {
+        int32_t * simi = distances + k * i;
+        idx_t * idxi = labels + k * i;
+
+        heap_heapify<HeapForL2> (k, simi, idxi);
+        KnnSearchResults res = {k, simi, idxi};
+        const uint8_t *q = x + i * code_size;
+
+        search_1_query_multihash (*this, q, res, n0, nlist, ndis);
+
+    }
+    indexBinaryHash_stats.nq += n;
+    indexBinaryHash_stats.n0 += n0;
+    indexBinaryHash_stats.nlist += nlist;
+    indexBinaryHash_stats.ndis += ndis;
+}
+
+size_t IndexBinaryMultiHash::hashtable_size() const
+{
+    size_t tot = 0;
+    for (auto map: maps) {
+        tot += map.size();
+    }
+
+    return tot;
+}
+
+
+}