faiss 0.6.0 → 0.6.1

data/vendor/faiss/faiss/invlists/InvertedLists.cpp

@@ -21,8 +21,8 @@ InvertedListsIterator::~InvertedListsIterator() {}
  * InvertedLists implementation
  ******************************************/
 
-InvertedLists::InvertedLists(size_t nlist, size_t code_size)
-        : nlist(nlist), code_size(code_size) {}
+InvertedLists::InvertedLists(size_t nlist_, size_t code_size_)
+        : nlist(nlist_), code_size(code_size_) {}
 
 InvertedLists::~InvertedLists() {}
 
@@ -70,7 +70,7 @@ void InvertedLists::reset() {
 
 void InvertedLists::merge_from(InvertedLists* oivf, size_t add_id) {
 #pragma omp parallel for
-    for (idx_t i = 0; i < nlist; i++) {
+    for (idx_t i = 0; i < static_cast<idx_t>(nlist); i++) {
         size_t list_size = oivf->list_size(i);
         ScopedIds ids(oivf, i);
         if (add_id == 0) {
@@ -109,12 +109,12 @@ size_t InvertedLists::copy_subset_to(
         ntotal = compute_ntotal();
     }
 
-    for (idx_t list_no = 0; list_no < nlist; list_no++) {
+    for (size_t list_no = 0; list_no < nlist; list_no++) {
         size_t n = list_size(list_no);
         ScopedIds ids_in(this, list_no);
 
         if (subset_type == SUBSET_TYPE_ID_RANGE) {
-            for (idx_t i = 0; i < n; i++) {
+            for (size_t i = 0; i < n; i++) {
                 idx_t id = ids_in[i];
                 if (a1 <= id && id < a2) {
                     oivf.add_entry(
@@ -125,7 +125,7 @@ size_t InvertedLists::copy_subset_to(
                 }
             }
         } else if (subset_type == SUBSET_TYPE_ID_MOD) {
-            for (idx_t i = 0; i < n; i++) {
+            for (size_t i = 0; i < n; i++) {
                 idx_t id = ids_in[i];
                 if (id % a1 == a2) {
                     oivf.add_entry(
@@ -143,7 +143,7 @@ size_t InvertedLists::copy_subset_to(
             size_t next_accu_a2 = next_accu_n * a2 / ntotal;
             size_t i2 = next_accu_a2 - accu_a2;
 
-            for (idx_t i = i1; i < i2; i++) {
+            for (size_t i = i1; i < i2; i++) {
                 oivf.add_entry(
                         list_no,
                         get_single_id(list_no, i),
@@ -157,7 +157,7 @@ size_t InvertedLists::copy_subset_to(
             size_t i1 = n * a2 / a1;
             size_t i2 = n * (a2 + 1) / a1;
 
-            for (idx_t i = i1; i < i2; i++) {
+            for (size_t i = i1; i < i2; i++) {
                 oivf.add_entry(
                         list_no,
                         get_single_id(list_no, i),
@@ -166,7 +166,8 @@ size_t InvertedLists::copy_subset_to(
 
             n_added += i2 - i1;
         } else if (subset_type == SUBSET_TYPE_INVLIST) {
-            if (list_no >= a1 && list_no < a2) {
+            if (static_cast<idx_t>(list_no) >= a1 &&
+                static_cast<idx_t>(list_no) < a2) {
                 oivf.add_entries(
                         list_no,
                         n,
@@ -269,10 +270,10 @@ InvertedListsIterator* InvertedLists::get_iterator(
  * ArrayInvertedLists implementation
  ******************************************/
 
-ArrayInvertedLists::ArrayInvertedLists(size_t nlist, size_t code_size)
-        : InvertedLists(nlist, code_size) {
-    ids.resize(nlist);
-    codes.resize(nlist);
+ArrayInvertedLists::ArrayInvertedLists(size_t nlist_in, size_t code_size_in)
+        : InvertedLists(nlist_in, code_size_in) {
+    ids.resize(nlist_in);
+    codes.resize(nlist_in);
 }
 
 size_t ArrayInvertedLists::add_entries(
@@ -288,7 +289,9 @@ size_t ArrayInvertedLists::add_entries(
     ids[list_no].resize(o + n_entry);
     memcpy(&ids[list_no][o], ids_in, sizeof(ids_in[0]) * n_entry);
     codes[list_no].resize((o + n_entry) * code_size);
-    memcpy(&codes[list_no][o * code_size], code, code_size * n_entry);
+    if (code_size > 0) {
+        memcpy(&codes[list_no][o * code_size], code, code_size * n_entry);
+    }
     return o;
 }
 
@@ -327,7 +330,11 @@ void ArrayInvertedLists::update_entries(
     assert(list_no < nlist);
     assert(n_entry + offset <= ids[list_no].size());
     memcpy(&ids[list_no][offset], ids_in, sizeof(ids_in[0]) * n_entry);
-    memcpy(&codes[list_no][offset * code_size], codes_in, code_size * n_entry);
+    if (code_size > 0) {
+        memcpy(&codes[list_no][offset * code_size],
+               codes_in,
+               code_size * n_entry);
+    }
 }
 
 void ArrayInvertedLists::permute_invlists(const idx_t* map) {
@@ -351,23 +358,25 @@ ArrayInvertedLists::~ArrayInvertedLists() {}
  **********************************************/
 
 ArrayInvertedListsPanorama::ArrayInvertedListsPanorama(
-        size_t nlist,
-        size_t code_size,
-        size_t n_levels)
-        : ArrayInvertedLists(nlist, code_size),
-          n_levels(n_levels),
+        size_t nlist_in,
+        size_t code_size_in,
+        size_t n_levels_in,
+        size_t batch_size)
+        : ArrayInvertedLists(nlist_in, code_size_in),
+          n_levels(n_levels_in),
           level_width(
-                  (((code_size / sizeof(float)) + n_levels - 1) / n_levels) *
+                  (((code_size_in / sizeof(float)) + n_levels_in - 1) /
+                   n_levels_in) *
                   sizeof(float)),
-          pano(code_size, n_levels, kBatchSize) {
-    FAISS_THROW_IF_NOT(n_levels > 0);
-    FAISS_THROW_IF_NOT(code_size % sizeof(float) == 0);
+          pano(code_size_in, n_levels_in, batch_size) {
+    FAISS_THROW_IF_NOT(n_levels_in > 0);
+    FAISS_THROW_IF_NOT(code_size_in % sizeof(float) == 0);
     FAISS_THROW_IF_NOT_MSG(
             !use_iterator,
             "IndexIVFFlatPanorama does not support iterators, use vanilla IndexIVFFlat instead");
     FAISS_ASSERT(level_width % sizeof(float) == 0);
 
-    cum_sums.resize(nlist);
+    cum_sums.resize(nlist_in);
 }
 
 const float* ArrayInvertedListsPanorama::get_cum_sums(size_t list_no) const {
@@ -387,9 +396,9 @@ size_t ArrayInvertedListsPanorama::add_entries(
     memcpy(&ids[list_no][o], ids_in, sizeof(ids_in[0]) * n_entry);
 
     size_t new_size = o + n_entry;
-    size_t num_batches = (new_size + kBatchSize - 1) / kBatchSize;
-    codes[list_no].resize(num_batches * kBatchSize * code_size);
-    cum_sums[list_no].resize(num_batches * kBatchSize * (n_levels + 1));
+    size_t num_batches = (new_size + pano.batch_size - 1) / pano.batch_size;
+    codes[list_no].resize(num_batches * pano.batch_size * code_size);
+    cum_sums[list_no].resize(num_batches * pano.batch_size * (n_levels + 1));
 
     // Cast to float* is safe here as we guarantee codes are always float
     // vectors for `IndexIVFFlatPanorama` (verified by the constructor).
@@ -423,9 +432,9 @@ void ArrayInvertedListsPanorama::update_entries(
 void ArrayInvertedListsPanorama::resize(size_t list_no, size_t new_size) {
     ids[list_no].resize(new_size);
 
-    size_t num_batches = (new_size + kBatchSize - 1) / kBatchSize;
-    codes[list_no].resize(num_batches * kBatchSize * code_size);
-    cum_sums[list_no].resize(num_batches * kBatchSize * (n_levels + 1));
+    size_t num_batches = (new_size + pano.batch_size - 1) / pano.batch_size;
+    codes[list_no].resize(num_batches * pano.batch_size * code_size);
+    cum_sums[list_no].resize(num_batches * pano.batch_size * (n_levels + 1));
 }
 
 const uint8_t* ArrayInvertedListsPanorama::get_single_code(
@@ -444,11 +453,11 @@ const uint8_t* ArrayInvertedListsPanorama::get_single_code(
 
 void ArrayInvertedListsPanorama::release_codes(
         size_t list_no,
-        const uint8_t* codes) const {
+        const uint8_t* codes_in) const {
     // Only delete if it's heap-allocated (from get_single_code).
     // If it's from get_codes (raw storage), it will be codes[list_no].data()
-    if (codes != this->codes[list_no].data()) {
-        delete[] codes;
+    if (codes_in != this->codes[list_no].data()) {
+        delete[] codes_in;
     }
 }
 
@@ -503,7 +512,7 @@ HStackInvertedLists::HStackInvertedLists(int nil, const InvertedLists** ils_in)
 
 size_t HStackInvertedLists::list_size(size_t list_no) const {
     size_t sz = 0;
-    for (int i = 0; i < ils.size(); i++) {
+    for (size_t i = 0; i < ils.size(); i++) {
         const InvertedLists* il = ils[i];
         sz += il->list_size(list_no);
     }
@@ -513,7 +522,7 @@ size_t HStackInvertedLists::list_size(size_t list_no) const {
 const uint8_t* HStackInvertedLists::get_codes(size_t list_no) const {
     uint8_t *codes = new uint8_t[code_size * list_size(list_no)], *c = codes;
 
-    for (int i = 0; i < ils.size(); i++) {
+    for (size_t i = 0; i < ils.size(); i++) {
         const InvertedLists* il = ils[i];
         size_t sz = il->list_size(list_no) * code_size;
         if (sz > 0) {
@@ -527,7 +536,7 @@ const uint8_t* HStackInvertedLists::get_codes(size_t list_no) const {
 const uint8_t* HStackInvertedLists::get_single_code(
         size_t list_no,
         size_t offset) const {
-    for (int i = 0; i < ils.size(); i++) {
+    for (size_t i = 0; i < ils.size(); i++) {
         const InvertedLists* il = ils[i];
         size_t sz = il->list_size(list_no);
         if (offset < sz) {
@@ -548,7 +557,7 @@ void HStackInvertedLists::release_codes(size_t, const uint8_t* codes) const {
 const idx_t* HStackInvertedLists::get_ids(size_t list_no) const {
     idx_t *ids = new idx_t[list_size(list_no)], *c = ids;
 
-    for (int i = 0; i < ils.size(); i++) {
+    for (size_t i = 0; i < ils.size(); i++) {
         const InvertedLists* il = ils[i];
         size_t sz = il->list_size(list_no);
         if (sz > 0) {
@@ -560,7 +569,7 @@ const idx_t* HStackInvertedLists::get_ids(size_t list_no) const {
 }
 
 idx_t HStackInvertedLists::get_single_id(size_t list_no, size_t offset) const {
-    for (int i = 0; i < ils.size(); i++) {
+    for (size_t i = 0; i < ils.size(); i++) {
         const InvertedLists* il = ils[i];
         size_t sz = il->list_size(list_no);
         if (offset < sz) {
@@ -575,11 +584,11 @@ void HStackInvertedLists::release_ids(size_t, const idx_t* ids) const {
     delete[] ids;
 }
 
-void HStackInvertedLists::prefetch_lists(const idx_t* list_nos, int nlist)
+void HStackInvertedLists::prefetch_lists(const idx_t* list_nos, int nlist_in)
         const {
-    for (int i = 0; i < ils.size(); i++) {
+    for (size_t i = 0; i < ils.size(); i++) {
         const InvertedLists* il = ils[i];
-        il->prefetch_lists(list_nos, nlist);
+        il->prefetch_lists(list_nos, nlist_in);
     }
 }
 
@@ -590,20 +599,21 @@ void HStackInvertedLists::prefetch_lists(const idx_t* list_nos, int nlist)
 namespace {
 
 idx_t translate_list_no(const SliceInvertedLists* sil, idx_t list_no) {
-    FAISS_THROW_IF_NOT(list_no >= 0 && list_no < sil->nlist);
+    FAISS_THROW_IF_NOT(
+            list_no >= 0 && static_cast<size_t>(list_no) < sil->nlist);
     return list_no + sil->i0;
 }
 
 } // namespace
 
 SliceInvertedLists::SliceInvertedLists(
-        const InvertedLists* il,
-        idx_t i0,
-        idx_t i1)
-        : ReadOnlyInvertedLists(i1 - i0, il->code_size),
-          il(il),
-          i0(i0),
-          i1(i1) {}
+        const InvertedLists* il_,
+        idx_t i0_,
+        idx_t i1_)
+        : ReadOnlyInvertedLists(i1_ - i0_, il_->code_size),
+          il(il_),
+          i0(i0_),
+          i1(i1_) {}
 
 size_t SliceInvertedLists::list_size(size_t list_no) const {
     return il->list_size(translate_list_no(this, list_no));
@@ -636,10 +646,10 @@ void SliceInvertedLists::release_ids(size_t list_no, const idx_t* ids) const {
     return il->release_ids(translate_list_no(this, list_no), ids);
 }
 
-void SliceInvertedLists::prefetch_lists(const idx_t* list_nos, int nlist)
+void SliceInvertedLists::prefetch_lists(const idx_t* list_nos, int nlist_in)
         const {
     std::vector<idx_t> translated_list_nos;
-    for (int j = 0; j < nlist; j++) {
+    for (size_t j = 0; j < static_cast<size_t>(nlist_in); j++) {
         idx_t list_no = list_nos[j];
         if (list_no < 0) {
             continue;
@@ -657,7 +667,8 @@ namespace {
 
 // find the invlist this number belongs to
 int translate_list_no(const VStackInvertedLists* vil, idx_t list_no) {
-    FAISS_THROW_IF_NOT(list_no >= 0 && list_no < vil->nlist);
+    FAISS_THROW_IF_NOT(
+            list_no >= 0 && static_cast<size_t>(list_no) < vil->nlist);
     int i0 = 0, i1 = vil->ils.size();
     const idx_t* cumsz = vil->cumsz.data();
     while (i0 + 1 < i1) {
@@ -740,11 +751,11 @@ void VStackInvertedLists::release_ids(size_t list_no, const idx_t* ids) const {
     return ils[i]->release_ids(list_no, ids);
 }
 
-void VStackInvertedLists::prefetch_lists(const idx_t* list_nos, int nlist)
+void VStackInvertedLists::prefetch_lists(const idx_t* list_nos, int nlist_in)
         const {
-    std::vector<int> ilno(nlist, -1);
+    std::vector<int> ilno(nlist_in, -1);
     std::vector<int> n_per_il(ils.size(), 0);
-    for (int j = 0; j < nlist; j++) {
+    for (int j = 0; j < nlist_in; j++) {
         idx_t list_no = list_nos[j];
         if (list_no < 0) {
             continue;
@@ -753,11 +764,11 @@ void VStackInvertedLists::prefetch_lists(const idx_t* list_nos, int nlist)
         n_per_il[i]++;
     }
     std::vector<int> cum_n_per_il(ils.size() + 1, 0);
-    for (int j = 0; j < ils.size(); j++) {
+    for (size_t j = 0; j < ils.size(); j++) {
         cum_n_per_il[j + 1] = cum_n_per_il[j] + n_per_il[j];
     }
     std::vector<idx_t> sorted_list_nos(cum_n_per_il.back());
-    for (int j = 0; j < nlist; j++) {
+    for (int j = 0; j < nlist_in; j++) {
         idx_t list_no = list_nos[j];
         if (list_no < 0) {
             continue;
@@ -768,7 +779,7 @@ void VStackInvertedLists::prefetch_lists(const idx_t* list_nos, int nlist)
     }
 
     int i0 = 0;
-    for (int j = 0; j < ils.size(); j++) {
+    for (size_t j = 0; j < ils.size(); j++) {
         int i1 = i0 + n_per_il[j];
         if (i1 > i0) {
             ils[j]->prefetch_lists(sorted_list_nos.data() + i0, i1 - i0);
@@ -782,11 +793,11 @@ void VStackInvertedLists::prefetch_lists(const idx_t* list_nos, int nlist)
  ******************************************/
 
 MaskedInvertedLists::MaskedInvertedLists(
-        const InvertedLists* il0,
-        const InvertedLists* il1)
-        : ReadOnlyInvertedLists(il0->nlist, il0->code_size),
-          il0(il0),
-          il1(il1) {
+        const InvertedLists* il0_in,
+        const InvertedLists* il1_in)
+        : ReadOnlyInvertedLists(il0_in->nlist, il0_in->code_size),
+          il0(il0_in),
+          il1(il1_in) {
     FAISS_THROW_IF_NOT(il1->nlist == nlist);
     FAISS_THROW_IF_NOT(il1->code_size == code_size);
 }
@@ -829,10 +840,10 @@ const uint8_t* MaskedInvertedLists::get_single_code(
     return (sz ? il0 : il1)->get_single_code(list_no, offset);
 }
 
-void MaskedInvertedLists::prefetch_lists(const idx_t* list_nos, int nlist)
+void MaskedInvertedLists::prefetch_lists(const idx_t* list_nos, int nlist_in)
         const {
     std::vector<idx_t> list0, list1;
-    for (int i = 0; i < nlist; i++) {
+    for (int i = 0; i < nlist_in; i++) {
         idx_t list_no = list_nos[i];
         if (list_no < 0) {
             continue;
@@ -849,11 +860,11 @@ void MaskedInvertedLists::prefetch_lists(const idx_t* list_nos, int nlist)
  ******************************************/
 
 StopWordsInvertedLists::StopWordsInvertedLists(
-        const InvertedLists* il0,
-        size_t maxsize)
-        : ReadOnlyInvertedLists(il0->nlist, il0->code_size),
-          il0(il0),
-          maxsize(maxsize) {}
+        const InvertedLists* il0_in,
+        size_t maxsize_in)
+        : ReadOnlyInvertedLists(il0_in->nlist, il0_in->code_size),
+          il0(il0_in),
+          maxsize(maxsize_in) {}
 
 size_t StopWordsInvertedLists::list_size(size_t list_no) const {
     size_t sz = il0->list_size(list_no);
@@ -896,10 +907,10 @@ const uint8_t* StopWordsInvertedLists::get_single_code(
     return il0->get_single_code(list_no, offset);
 }
 
-void StopWordsInvertedLists::prefetch_lists(const idx_t* list_nos, int nlist)
+void StopWordsInvertedLists::prefetch_lists(const idx_t* list_nos, int nlist_in)
         const {
     std::vector<idx_t> list0;
-    for (int i = 0; i < nlist; i++) {
+    for (int i = 0; i < nlist_in; i++) {
         idx_t list_no = list_nos[i];
         if (list_no < 0) {
             continue;
@@ -911,4 +922,77 @@ void StopWordsInvertedLists::prefetch_lists(const idx_t* list_nos, int nlist)
     il0->prefetch_lists(list0.data(), list0.size());
 }
 
+/*****************************************
+ * CappedInvertedLists implementation
+ ******************************************/
+
+CappedInvertedLists::CappedInvertedLists(InvertedLists* il0, size_t maxsize)
+        : InvertedLists(il0->nlist, il0->code_size),
+          il0(il0),
+          maxsize(maxsize) {}
+
+size_t CappedInvertedLists::list_size(size_t list_no) const {
+    size_t sz = il0->list_size(list_no);
+    return sz < maxsize ? sz : maxsize;
+}
+
+size_t CappedInvertedLists::real_list_size(size_t list_no) const {
+    return il0->list_size(list_no);
+}
+
+const uint8_t* CappedInvertedLists::get_codes(size_t list_no) const {
+    return il0->get_codes(list_no);
+}
+
+const idx_t* CappedInvertedLists::get_ids(size_t list_no) const {
+    return il0->get_ids(list_no);
+}
+
+void CappedInvertedLists::release_codes(size_t list_no, const uint8_t* codes)
+        const {
+    il0->release_codes(list_no, codes);
+}
+
+void CappedInvertedLists::release_ids(size_t list_no, const idx_t* ids) const {
+    il0->release_ids(list_no, ids);
+}
+
+idx_t CappedInvertedLists::get_single_id(size_t list_no, size_t offset) const {
+    FAISS_THROW_IF_NOT(offset < list_size(list_no));
+    return il0->get_single_id(list_no, offset);
+}
+
+const uint8_t* CappedInvertedLists::get_single_code(
+        size_t list_no,
+        size_t offset) const {
+    FAISS_THROW_IF_NOT(offset < list_size(list_no));
+    return il0->get_single_code(list_no, offset);
+}
+
+void CappedInvertedLists::prefetch_lists(const idx_t* list_nos, int nlist)
+        const {
+    il0->prefetch_lists(list_nos, nlist);
+}
+
+size_t CappedInvertedLists::add_entries(
+        size_t list_no,
+        size_t n_entry,
+        const idx_t* ids,
+        const uint8_t* code) {
+    return il0->add_entries(list_no, n_entry, ids, code);
+}
+
+void CappedInvertedLists::update_entries(
+        size_t list_no,
+        size_t offset,
+        size_t n_entry,
+        const idx_t* ids,
+        const uint8_t* code) {
+    il0->update_entries(list_no, offset, n_entry, ids, code);
+}
+
+void CappedInvertedLists::resize(size_t list_no, size_t new_size) {
+    il0->resize(list_no, new_size);
+}
+
 } // namespace faiss