faiss 0.1.7 → 0.2.3

data/vendor/faiss/faiss/impl/AuxIndexStructures.cpp

@@ -14,18 +14,16 @@
 
 #include <faiss/impl/FaissAssert.h>
 
-
 namespace faiss {
 
-
 /***********************************************************************
  * RangeSearchResult
  ***********************************************************************/
 
-RangeSearchResult::RangeSearchResult (idx_t nq, bool alloc_lims): nq (nq) {
+RangeSearchResult::RangeSearchResult(idx_t nq, bool alloc_lims) : nq(nq) {
     if (alloc_lims) {
-        lims = new size_t [nq + 1];
-        memset (lims, 0, sizeof(*lims) * (nq + 1));
+        lims = new size_t[nq + 1];
+        memset(lims, 0, sizeof(*lims) * (nq + 1));
     } else {
         lims = nullptr;
     }
@@ -36,145 +34,129 @@ RangeSearchResult::RangeSearchResult (idx_t nq, bool alloc_lims): nq (nq) {
 
 /// called when lims contains the nb of elements result entries
 /// for each query
-void RangeSearchResult::do_allocation () {
+void RangeSearchResult::do_allocation() {
+    // works only if all the partial results are aggregated
+    // simulatenously
+    FAISS_THROW_IF_NOT(labels == nullptr && distances == nullptr);
     size_t ofs = 0;
     for (int i = 0; i < nq; i++) {
         size_t n = lims[i];
-        lims [i] = ofs;
+        lims[i] = ofs;
         ofs += n;
     }
-    lims [nq] = ofs;
-    labels = new idx_t [ofs];
-    distances = new float [ofs];
+    lims[nq] = ofs;
+    labels = new idx_t[ofs];
+    distances = new float[ofs];
 }
 
-RangeSearchResult::~RangeSearchResult () {
-    delete [] labels;
-    delete [] distances;
-    delete [] lims;
+RangeSearchResult::~RangeSearchResult() {
+    delete[] labels;
+    delete[] distances;
+    delete[] lims;
 }
 
-
-
-
-
 /***********************************************************************
  * BufferList
  ***********************************************************************/
 
-
-BufferList::BufferList (size_t buffer_size):
-    buffer_size (buffer_size)
-{
+BufferList::BufferList(size_t buffer_size) : buffer_size(buffer_size) {
     wp = buffer_size;
 }
 
-BufferList::~BufferList ()
-{
+BufferList::~BufferList() {
     for (int i = 0; i < buffers.size(); i++) {
-        delete [] buffers[i].ids;
-        delete [] buffers[i].dis;
+        delete[] buffers[i].ids;
+        delete[] buffers[i].dis;
     }
 }
 
-void BufferList::add (idx_t id, float dis) {
+void BufferList::add(idx_t id, float dis) {
     if (wp == buffer_size) { // need new buffer
         append_buffer();
     }
-    Buffer & buf = buffers.back();
-    buf.ids [wp] = id;
-    buf.dis [wp] = dis;
+    Buffer& buf = buffers.back();
+    buf.ids[wp] = id;
+    buf.dis[wp] = dis;
     wp++;
 }
 
-
-void BufferList::append_buffer ()
-{
-    Buffer buf = {new idx_t [buffer_size], new float [buffer_size]};
-    buffers.push_back (buf);
+void BufferList::append_buffer() {
+    Buffer buf = {new idx_t[buffer_size], new float[buffer_size]};
+    buffers.push_back(buf);
     wp = 0;
 }
 
 /// copy elemnts ofs:ofs+n-1 seen as linear data in the buffers to
 /// tables dest_ids, dest_dis
-void BufferList::copy_range (size_t ofs, size_t n,
-                             idx_t * dest_ids, float *dest_dis)
-{
+void BufferList::copy_range(
+        size_t ofs,
+        size_t n,
+        idx_t* dest_ids,
+        float* dest_dis) {
     size_t bno = ofs / buffer_size;
     ofs -= bno * buffer_size;
     while (n > 0) {
         size_t ncopy = ofs + n < buffer_size ? n : buffer_size - ofs;
-        Buffer buf = buffers [bno];
-        memcpy (dest_ids, buf.ids + ofs, ncopy * sizeof(*dest_ids));
-        memcpy (dest_dis, buf.dis + ofs, ncopy * sizeof(*dest_dis));
+        Buffer buf = buffers[bno];
+        memcpy(dest_ids, buf.ids + ofs, ncopy * sizeof(*dest_ids));
+        memcpy(dest_dis, buf.dis + ofs, ncopy * sizeof(*dest_dis));
         dest_ids += ncopy;
         dest_dis += ncopy;
         ofs = 0;
-        bno ++;
+        bno++;
         n -= ncopy;
     }
 }
 
-
 /***********************************************************************
  * RangeSearchPartialResult
  ***********************************************************************/
 
-void RangeQueryResult::add (float dis, idx_t id) {
+void RangeQueryResult::add(float dis, idx_t id) {
     nres++;
-    pres->add (id, dis);
+    pres->add(id, dis);
 }
 
-
-
-RangeSearchPartialResult::RangeSearchPartialResult (RangeSearchResult * res_in):
-    BufferList(res_in->buffer_size),
-    res(res_in)
-{}
-
+RangeSearchPartialResult::RangeSearchPartialResult(RangeSearchResult* res_in)
+        : BufferList(res_in->buffer_size), res(res_in) {}
 
 /// begin a new result
-RangeQueryResult &
-    RangeSearchPartialResult::new_result (idx_t qno)
-{
+RangeQueryResult& RangeSearchPartialResult::new_result(idx_t qno) {
     RangeQueryResult qres = {qno, 0, this};
-    queries.push_back (qres);
+    queries.push_back(qres);
     return queries.back();
 }
 
-
-void RangeSearchPartialResult::finalize ()
-{
-    set_lims ();
+void RangeSearchPartialResult::finalize() {
+    set_lims();
 #pragma omp barrier
 
 #pragma omp single
-    res->do_allocation ();
+    res->do_allocation();
 
 #pragma omp barrier
-    copy_result ();
+    copy_result();
 }
 
-
 /// called by range_search before do_allocation
-void RangeSearchPartialResult::set_lims ()
-{
+void RangeSearchPartialResult::set_lims() {
     for (int i = 0; i < queries.size(); i++) {
-        RangeQueryResult & qres = queries[i];
+        RangeQueryResult& qres = queries[i];
         res->lims[qres.qno] = qres.nres;
     }
 }
 
 /// called by range_search after do_allocation
-void RangeSearchPartialResult::copy_result (bool incremental)
-{
+void RangeSearchPartialResult::copy_result(bool incremental) {
     size_t ofs = 0;
     for (int i = 0; i < queries.size(); i++) {
-        RangeQueryResult & qres = queries[i];
+        RangeQueryResult& qres = queries[i];
 
-        copy_range (ofs, qres.nres,
-                    res->labels + res->lims[qres.qno],
-                    res->distances + res->lims[qres.qno]);
+        copy_range(
+                ofs,
+                qres.nres,
+                res->labels + res->lims[qres.qno],
+                res->distances + res->lims[qres.qno]);
         if (incremental) {
             res->lims[qres.qno] += qres.nres;
         }
@@ -182,26 +164,28 @@ void RangeSearchPartialResult::copy_result (bool incremental)
     }
 }
 
-void RangeSearchPartialResult::merge (std::vector <RangeSearchPartialResult *> &
-                                      partial_results, bool do_delete)
-{
-
+void RangeSearchPartialResult::merge(
+        std::vector<RangeSearchPartialResult*>& partial_results,
+        bool do_delete) {
     int npres = partial_results.size();
-    if (npres == 0) return;
-    RangeSearchResult *result = partial_results[0]->res;
+    if (npres == 0)
+        return;
+    RangeSearchResult* result = partial_results[0]->res;
     size_t nx = result->nq;
 
     // count
-    for (const RangeSearchPartialResult * pres : partial_results) {
-        if (!pres) continue;
-        for (const RangeQueryResult &qres : pres->queries) {
+    for (const RangeSearchPartialResult* pres : partial_results) {
+        if (!pres)
+            continue;
+        for (const RangeQueryResult& qres : pres->queries) {
             result->lims[qres.qno] += qres.nres;
         }
     }
-    result->do_allocation ();
+    result->do_allocation();
     for (int j = 0; j < npres; j++) {
-        if (!partial_results[j]) continue;
-        partial_results[j]->copy_result (true);
+        if (!partial_results[j])
+            continue;
+        partial_results[j]->copy_result(true);
         if (do_delete) {
             delete partial_results[j];
             partial_results[j] = nullptr;
@@ -210,22 +194,19 @@ void RangeSearchPartialResult::merge (std::vector <RangeSearchPartialResult *> &
 
     // reset the limits
     for (size_t i = nx; i > 0; i--) {
-        result->lims [i] = result->lims [i - 1];
+        result->lims[i] = result->lims[i - 1];
     }
-    result->lims [0] = 0;
+    result->lims[0] = 0;
 }
 
 /***********************************************************************
  * IDSelectorRange
  ***********************************************************************/
 
-IDSelectorRange::IDSelectorRange (idx_t imin, idx_t imax):
-    imin (imin), imax (imax)
-{
-}
+IDSelectorRange::IDSelectorRange(idx_t imin, idx_t imax)
+        : imin(imin), imax(imax) {}
 
-bool IDSelectorRange::is_member (idx_t id) const
-{
+bool IDSelectorRange::is_member(idx_t id) const {
     return id >= imin && id < imax;
 }
 
@@ -233,33 +214,29 @@ bool IDSelectorRange::is_member (idx_t id) const
  * IDSelectorArray
  ***********************************************************************/
 
-IDSelectorArray::IDSelectorArray (size_t n, const idx_t *ids):
-   n (n), ids(ids)
-{
-}
+IDSelectorArray::IDSelectorArray(size_t n, const idx_t* ids) : n(n), ids(ids) {}
 
-bool IDSelectorArray::is_member (idx_t id) const
-{
+bool IDSelectorArray::is_member(idx_t id) const {
     for (idx_t i = 0; i < n; i++) {
-        if (ids[i] == id) return true;
+        if (ids[i] == id)
+            return true;
     }
     return false;
 }
 
-
 /***********************************************************************
  * IDSelectorBatch
  ***********************************************************************/
 
-IDSelectorBatch::IDSelectorBatch (size_t n, const idx_t *indices)
-{
+IDSelectorBatch::IDSelectorBatch(size_t n, const idx_t* indices) {
     nbits = 0;
-    while (n > (1L << nbits)) nbits++;
+    while (n > (1L << nbits))
+        nbits++;
     nbits += 5;
     // for n = 1M, nbits = 25 is optimal, see P56659518
 
     mask = (1L << nbits) - 1;
-    bloom.resize (1UL << (nbits - 3), 0);
+    bloom.resize(1UL << (nbits - 3), 0);
     for (long i = 0; i < n; i++) {
         Index::idx_t id = indices[i];
         set.insert(id);
@@ -268,39 +245,36 @@ IDSelectorBatch::IDSelectorBatch (size_t n, const idx_t *indices)
     }
 }
 
-bool IDSelectorBatch::is_member (idx_t i) const
-{
+bool IDSelectorBatch::is_member(idx_t i) const {
     long im = i & mask;
-    if(!(bloom[im>>3] & (1 << (im & 7)))) {
+    if (!(bloom[im >> 3] & (1 << (im & 7)))) {
         return 0;
     }
     return set.count(i);
 }
 
-
 /***********************************************************
  * Interrupt callback
  ***********************************************************/
 
-
 std::unique_ptr<InterruptCallback> InterruptCallback::instance;
 
 std::mutex InterruptCallback::lock;
 
-void InterruptCallback::clear_instance () {
-    delete instance.release ();
+void InterruptCallback::clear_instance() {
+    delete instance.release();
 }
 
-void InterruptCallback::check () {
+void InterruptCallback::check() {
     if (!instance.get()) {
         return;
     }
-    if (instance->want_interrupt ()) {
-        FAISS_THROW_MSG ("computation interrupted");
+    if (instance->want_interrupt()) {
+        FAISS_THROW_MSG("computation interrupted");
     }
 }
 
-bool InterruptCallback::is_interrupted () {
+bool InterruptCallback::is_interrupted() {
     if (!instance.get()) {
         return false;
     }
@@ -308,8 +282,7 @@ bool InterruptCallback::is_interrupted () {
     return instance->want_interrupt();
 }
 
-
-size_t InterruptCallback::get_period_hint (size_t flops) {
+size_t InterruptCallback::get_period_hint(size_t flops) {
     if (!instance.get()) {
         return 1L << 30; // never check
     }
@@ -317,7 +290,4 @@ size_t InterruptCallback::get_period_hint (size_t flops) {
     return std::max((size_t)10 * 10 * 1000 * 1000 / (flops + 1), (size_t)1);
 }
 
-
-
-
 } // namespace faiss

data/vendor/faiss/faiss/impl/AuxIndexStructures.h

@@ -15,15 +15,15 @@
 
 #include <stdint.h>
 
-#include <vector>
-#include <unordered_set>
+#include <cstring>
 #include <memory>
 #include <mutex>
+#include <unordered_set>
+#include <vector>
 
 #include <faiss/Index.h>
 #include <faiss/impl/platform_macros.h>
 
-
 namespace faiss {
 
 /** The objective is to have a simple result structure while
@@ -31,42 +31,39 @@ namespace faiss {
  *  do_allocation can be overloaded to allocate the result tables in
  *  the matrix type of a scripting language like Lua or Python. */
 struct RangeSearchResult {
-    size_t nq;      ///< nb of queries
-    size_t *lims;   ///< size (nq + 1)
+    size_t nq;    ///< nb of queries
+    size_t* lims; ///< size (nq + 1)
 
     typedef Index::idx_t idx_t;
 
-    idx_t *labels;     ///< result for query i is labels[lims[i]:lims[i+1]]
-    float *distances;  ///< corresponding distances (not sorted)
+    idx_t* labels;    ///< result for query i is labels[lims[i]:lims[i+1]]
+    float* distances; ///< corresponding distances (not sorted)
 
     size_t buffer_size; ///< size of the result buffers used
 
     /// lims must be allocated on input to range_search.
-    explicit RangeSearchResult (idx_t nq, bool alloc_lims=true);
+    explicit RangeSearchResult(idx_t nq, bool alloc_lims = true);
 
     /// called when lims contains the nb of elements result entries
     /// for each query
 
-    virtual void do_allocation ();
+    virtual void do_allocation();
 
-    virtual ~RangeSearchResult ();
+    virtual ~RangeSearchResult();
 };
 
-
 /** Encapsulates a set of ids to remove. */
 struct IDSelector {
     typedef Index::idx_t idx_t;
-    virtual bool is_member (idx_t id) const = 0;
+    virtual bool is_member(idx_t id) const = 0;
     virtual ~IDSelector() {}
 };
 
-
-
 /** remove ids between [imni, imax) */
-struct IDSelectorRange: IDSelector {
+struct IDSelectorRange : IDSelector {
     idx_t imin, imax;
 
-    IDSelectorRange (idx_t imin, idx_t imax);
+    IDSelectorRange(idx_t imin, idx_t imax);
     bool is_member(idx_t id) const override;
     ~IDSelectorRange() override {}
 };
@@ -76,11 +73,11 @@ struct IDSelectorRange: IDSelector {
  * this is inefficient in most cases, except for IndexIVF with
  * maintain_direct_map
  */
-struct IDSelectorArray: IDSelector {
+struct IDSelectorArray : IDSelector {
     size_t n;
-    const idx_t *ids;
+    const idx_t* ids;
 
-    IDSelectorArray (size_t n, const idx_t *ids);
+    IDSelectorArray(size_t n, const idx_t* ids);
     bool is_member(idx_t id) const override;
     ~IDSelectorArray() override {}
 };
@@ -91,8 +88,7 @@ struct IDSelectorArray: IDSelector {
  * unordered_set are just the least significant bits of the id. This
  * works fine for random ids or ids in sequences but will produce many
  * hash collisions if lsb's are always the same */
-struct IDSelectorBatch: IDSelector {
-
+struct IDSelectorBatch : IDSelector {
     std::unordered_set<idx_t> set;
 
     typedef unsigned char uint8_t;
@@ -100,7 +96,7 @@ struct IDSelectorBatch: IDSelector {
     int nbits;
     idx_t mask;
 
-    IDSelectorBatch (size_t n, const idx_t *indices);
+    IDSelectorBatch(size_t n, const idx_t* indices);
     bool is_member(idx_t id) const override;
     ~IDSelectorBatch() override {}
 };
@@ -124,28 +120,26 @@ struct BufferList {
     size_t buffer_size;
 
     struct Buffer {
-        idx_t *ids;
-        float *dis;
+        idx_t* ids;
+        float* dis;
     };
 
     std::vector<Buffer> buffers;
     size_t wp; ///< write pointer in the last buffer.
 
-    explicit BufferList (size_t buffer_size);
+    explicit BufferList(size_t buffer_size);
 
-    ~BufferList ();
+    ~BufferList();
 
     /// create a new buffer
-    void append_buffer ();
+    void append_buffer();
 
     /// add one result, possibly appending a new buffer if needed
-    void add (idx_t id, float dis);
+    void add(idx_t id, float dis);
 
     /// copy elemnts ofs:ofs+n-1 seen as linear data in the buffers to
     /// tables dest_ids, dest_dis
-    void copy_range (size_t ofs, size_t n,
-                     idx_t * dest_ids, float *dest_dis);
-
+    void copy_range(size_t ofs, size_t n, idx_t* dest_ids, float* dest_dis);
 };
 
 struct RangeSearchPartialResult;
@@ -153,46 +147,45 @@ struct RangeSearchPartialResult;
 /// result structure for a single query
 struct RangeQueryResult {
     using idx_t = Index::idx_t;
-    idx_t qno;    //< id of the query
-    size_t nres;  //< nb of results for this query
-    RangeSearchPartialResult * pres;
+    idx_t qno;   //< id of the query
+    size_t nres; //< nb of results for this query
+    RangeSearchPartialResult* pres;
 
     /// called by search function to report a new result
-    void add (float dis, idx_t id);
+    void add(float dis, idx_t id);
 };
 
 /// the entries in the buffers are split per query
-struct RangeSearchPartialResult: BufferList {
-    RangeSearchResult * res;
+struct RangeSearchPartialResult : BufferList {
+    RangeSearchResult* res;
 
     /// eventually the result will be stored in res_in
-    explicit RangeSearchPartialResult (RangeSearchResult * res_in);
+    explicit RangeSearchPartialResult(RangeSearchResult* res_in);
 
     /// query ids + nb of results per query.
     std::vector<RangeQueryResult> queries;
 
     /// begin a new result
-    RangeQueryResult & new_result (idx_t qno);
+    RangeQueryResult& new_result(idx_t qno);
 
     /*****************************************
      * functions used at the end of the search to merge the result
      * lists */
-    void finalize ();
+    void finalize();
 
     /// called by range_search before do_allocation
-    void set_lims ();
+    void set_lims();
 
     /// called by range_search after do_allocation
-    void copy_result (bool incremental = false);
+    void copy_result(bool incremental = false);
 
     /// merge a set of PartialResult's into one RangeSearchResult
     /// on ouptut the partialresults are empty!
-    static void merge (std::vector <RangeSearchPartialResult *> &
-                       partial_results, bool do_delete=true);
-
+    static void merge(
+            std::vector<RangeSearchPartialResult*>& partial_results,
+            bool do_delete = true);
 };
 
-
 /***********************************************************
  * The distance computer maintains a current query and computes
  * distances to elements in an index that supports random access.
@@ -202,19 +195,19 @@ struct RangeSearchPartialResult: BufferList {
  * instantiate one from each thread if needed.
  ***********************************************************/
 struct DistanceComputer {
-     using idx_t = Index::idx_t;
+    using idx_t = Index::idx_t;
 
-     /// called before computing distances. Pointer x should remain valid
-     /// while operator () is called
-     virtual void set_query(const float *x) = 0;
+    /// called before computing distances. Pointer x should remain valid
+    /// while operator () is called
+    virtual void set_query(const float* x) = 0;
 
-     /// compute distance of vector i to current query
-     virtual float operator () (idx_t i) = 0;
+    /// compute distance of vector i to current query
+    virtual float operator()(idx_t i) = 0;
 
-     /// compute distance between two stored vectors
-     virtual float symmetric_dis (idx_t i, idx_t j) = 0;
+    /// compute distance between two stored vectors
+    virtual float symmetric_dis(idx_t i, idx_t j) = 0;
 
-     virtual ~DistanceComputer() {}
+    virtual ~DistanceComputer() {}
 };
 
 /***********************************************************
@@ -222,7 +215,7 @@ struct DistanceComputer {
  ***********************************************************/
 
 struct FAISS_API InterruptCallback {
-    virtual bool want_interrupt () = 0;
+    virtual bool want_interrupt() = 0;
     virtual ~InterruptCallback() {}
 
     // lock that protects concurrent calls to is_interrupted
@@ -230,7 +223,7 @@ struct FAISS_API InterruptCallback {
 
     static std::unique_ptr<InterruptCallback> instance;
 
-    static void clear_instance ();
+    static void clear_instance();
 
     /** check if:
      * - an interrupt callback is set
@@ -238,23 +231,46 @@ struct FAISS_API InterruptCallback {
      * if this is the case, then throw an exception. Should not be called
      * from multiple threads.
      */
-    static void check ();
+    static void check();
 
     /// same as check() but return true if is interrupted instead of
     /// throwing. Can be called from multiple threads.
-    static bool is_interrupted ();
+    static bool is_interrupted();
 
     /** assuming each iteration takes a certain number of flops, what
      * is a reasonable interval to check for interrupts?
      */
-    static size_t get_period_hint (size_t flops);
-
+    static size_t get_period_hint(size_t flops);
 };
 
+/// set implementation optimized for fast access.
+struct VisitedTable {
+    std::vector<uint8_t> visited;
+    int visno;
+
+    explicit VisitedTable(int size) : visited(size), visno(1) {}
+
+    /// set flag #no to true
+    void set(int no) {
+        visited[no] = visno;
+    }
+
+    /// get flag #no
+    bool get(int no) const {
+        return visited[no] == visno;
+    }
+
+    /// reset all flags to false
+    void advance() {
+        visno++;
+        if (visno == 250) {
+            // 250 rather than 255 because sometimes we use visno and visno+1
+            memset(visited.data(), 0, sizeof(visited[0]) * visited.size());
+            visno = 1;
+        }
+    }
+};
 
-
-}; // namespace faiss
-
-
+} // namespace faiss
 
 #endif