faiss 0.6.0 → 0.6.1

data/vendor/faiss/faiss/IndexScalarQuantizer.cpp

@@ -26,10 +26,10 @@ namespace faiss {
  ********************************************************************/
 
 IndexScalarQuantizer::IndexScalarQuantizer(
-        int d,
+        int d_in,
         ScalarQuantizer::QuantizerType qtype,
         MetricType metric)
-        : IndexFlatCodes(0, d, metric), sq(d, qtype) {
+        : IndexFlatCodes(0, d_in, metric), sq(d_in, qtype) {
     is_trained = qtype == ScalarQuantizer::QT_fp16 ||
             qtype == ScalarQuantizer::QT_8bit_direct ||
             qtype == ScalarQuantizer::QT_bf16 ||
@@ -117,21 +117,26 @@ void IndexScalarQuantizer::sa_decode(idx_t n, const uint8_t* bytes, float* x)
  ********************************************************************/
 
 IndexIVFScalarQuantizer::IndexIVFScalarQuantizer(
-        Index* quantizer,
-        size_t d,
-        size_t nlist,
+        Index* quantizer_,
+        size_t d_,
+        size_t nlist_,
         ScalarQuantizer::QuantizerType qtype,
         MetricType metric,
-        bool by_residual,
-        bool own_invlists)
-        : IndexIVF(quantizer, d, nlist, 0, metric, own_invlists), sq(d, qtype) {
+        bool by_residual_,
+        bool own_invlists_)
+        : IndexIVF(quantizer_, d_, nlist_, 0, metric, own_invlists_),
+          sq(d_, qtype) {
     code_size = sq.code_size;
-    this->by_residual = by_residual;
+    this->by_residual = by_residual_;
     if (invlists) {
         // was not known at construction time
         invlists->code_size = code_size;
     }
     is_trained = false;
+    if (qtype == ScalarQuantizer::QT_0bit) {
+        by_residual = false;
+        is_trained = true; // no training needed
+    }
 }
 
 IndexIVFScalarQuantizer::IndexIVFScalarQuantizer() : IndexIVF() {
@@ -141,7 +146,7 @@ IndexIVFScalarQuantizer::IndexIVFScalarQuantizer() : IndexIVF() {
 void IndexIVFScalarQuantizer::train_encoder(
         idx_t n,
         const float* x,
-        const idx_t* assign) {
+        const idx_t* /*assign*/) {
     sq.train(n, x);
 }
 
@@ -155,6 +160,19 @@ void IndexIVFScalarQuantizer::encode_vectors(
         const idx_t* list_nos,
         uint8_t* codes,
         bool include_listnos) const {
+    if (sq.code_size == 0) {
+        // QT_0bit: nothing to encode, but handle coarse codes if needed
+        if (include_listnos) {
+            size_t coarse_size = coarse_code_size();
+            for (idx_t i = 0; i < n; i++) {
+                int64_t list_no = list_nos[i];
+                if (list_no >= 0) {
+                    encode_listno(list_no, codes + i * coarse_size);
+                }
+            }
+        }
+        return;
+    }
     std::unique_ptr<ScalarQuantizer::SQuantizer> squant(sq.select_quantizer());
     size_t coarse_size = include_listnos ? coarse_code_size() : 0;
     memset(codes, 0, (code_size + coarse_size) * n);
@@ -185,14 +203,42 @@ void IndexIVFScalarQuantizer::encode_vectors(
 void IndexIVFScalarQuantizer::decode_vectors(
         idx_t n,
         const uint8_t* codes,
-        const idx_t*,
+        const idx_t* list_nos,
         float* x) const {
+    if (sq.code_size == 0 && list_nos) {
+        // QT_0bit: reconstruct centroids if list_nos provided
+        for (idx_t i = 0; i < n; i++) {
+            quantizer->reconstruct(list_nos[i], x + i * d);
+        }
+        return;
+    }
     FAISS_THROW_IF_NOT(is_trained);
-    return sq.decode(codes, x, n);
+    sq.decode(codes, x, n);
+    if (by_residual) {
+        FAISS_THROW_IF_NOT_MSG(
+                list_nos, "decode_vectors with by_residual requires list_nos");
+#pragma omp parallel for if (n > 1000)
+        for (idx_t i = 0; i < n; i++) {
+            std::vector<float> centroid(d);
+            quantizer->reconstruct(list_nos[i], centroid.data());
+            for (size_t j = 0; j < static_cast<size_t>(d); j++) {
+                x[i * d + j] += centroid[j];
+            }
+        }
+    }
 }
 
 void IndexIVFScalarQuantizer::sa_decode(idx_t n, const uint8_t* codes, float* x)
         const {
+    if (sq.code_size == 0) {
+        size_t coarse_size = coarse_code_size();
+        for (idx_t i = 0; i < n; i++) {
+            const uint8_t* code = codes + i * coarse_size;
+            int64_t list_no = decode_listno(code);
+            quantizer->reconstruct(list_no, x + i * d);
+        }
+        return;
+    }
     std::unique_ptr<ScalarQuantizer::SQuantizer> squant(sq.select_quantizer());
     size_t coarse_size = coarse_code_size();
 
@@ -208,7 +254,7 @@ void IndexIVFScalarQuantizer::sa_decode(idx_t n, const uint8_t* codes, float* x)
             squant->decode_vector(code + coarse_size, xi);
             if (by_residual) {
                 quantizer->reconstruct(list_no, residual.data());
-                for (size_t j = 0; j < d; j++) {
+                for (size_t j = 0; j < static_cast<size_t>(d); j++) {
                     xi[j] += residual[j];
                 }
             }
@@ -223,6 +269,23 @@ void IndexIVFScalarQuantizer::add_core(
         const idx_t* coarse_idx,
         void* inverted_list_context) {
     FAISS_THROW_IF_NOT(is_trained);
+    if (sq.code_size == 0) {
+        // QT_0bit: just add IDs with empty codes
+        uint8_t dummy_code = 0;
+        DirectMapAdd dm_add(direct_map, n, xids);
+        for (idx_t i = 0; i < n; i++) {
+            int64_t list_no = coarse_idx[i];
+            if (list_no >= 0) {
+                int64_t id = xids ? xids[i] : ntotal + i;
+                size_t ofs = invlists->add_entry(list_no, id, &dummy_code);
+                dm_add.add(i, list_no, ofs);
+            } else {
+                dm_add.add(i, -1, 0);
+            }
+        }
+        ntotal += n;
+        return;
+    }
 
     std::unique_ptr<ScalarQuantizer::SQuantizer> squant(sq.select_quantizer());
 
@@ -236,7 +299,7 @@ void IndexIVFScalarQuantizer::add_core(
         int rank = omp_get_thread_num();
 
         // each thread takes care of a subset of lists
-        for (size_t i = 0; i < n; i++) {
+        for (idx_t i = 0; i < n; i++) {
             int64_t list_no = coarse_idx[i];
             if (list_no >= 0 && list_no % nt == rank) {
                 int64_t id = xids ? xids[i] : ntotal + i;
@@ -276,6 +339,11 @@ void IndexIVFScalarQuantizer::reconstruct_from_offset(
         int64_t list_no,
         int64_t offset,
         float* recons) const {
+    if (sq.code_size == 0) {
+        // QT_0bit: reconstruct from centroid
+        quantizer->reconstruct(list_no, recons);
+        return;
+    }
     const uint8_t* code = invlists->get_single_code(list_no, offset);
 
     if (by_residual) {

data/vendor/faiss/faiss/IndexShards.cpp

@@ -46,27 +46,28 @@ void translate_labels(int64_t n, idx_t* labels, int64_t translation) {
 
 template <typename IndexT>
 IndexShardsTemplate<IndexT>::IndexShardsTemplate(
-        idx_t d,
+        idx_t d_,
         bool threaded,
-        bool successive_ids)
-        : ThreadedIndex<IndexT>(d, threaded), successive_ids(successive_ids) {
+        bool successive_ids_)
+        : ThreadedIndex<IndexT>(static_cast<int>(d_), threaded),
+          successive_ids(successive_ids_) {
     sync_d(this);
 }
 
 template <typename IndexT>
 IndexShardsTemplate<IndexT>::IndexShardsTemplate(
-        int d,
+        int d_,
         bool threaded,
-        bool successive_ids)
-        : ThreadedIndex<IndexT>(d, threaded), successive_ids(successive_ids) {
+        bool successive_ids_)
+        : ThreadedIndex<IndexT>(d_, threaded), successive_ids(successive_ids_) {
     sync_d(this);
 }
 
 template <typename IndexT>
 IndexShardsTemplate<IndexT>::IndexShardsTemplate(
         bool threaded,
-        bool successive_ids)
-        : ThreadedIndex<IndexT>(threaded), successive_ids(successive_ids) {
+        bool successive_ids_)
+        : ThreadedIndex<IndexT>(threaded), successive_ids(successive_ids_) {
     sync_d(this);
 }
 
@@ -202,7 +203,7 @@ void IndexShardsTemplate<IndexT>::search(
         const SearchParameters* params) const {
     FAISS_THROW_IF_NOT(k > 0);
 
-    int64_t nshard = this->count();
+    int nshard = this->count();
 
     std::vector<distance_t> all_distances(nshard * k * n);
     std::vector<idx_t> all_labels(nshard * k * n);

data/vendor/faiss/faiss/IndexShardsIVF.cpp

@@ -43,13 +43,14 @@ void translate_labels(int64_t n, idx_t* labels, int64_t translation) {
  ************************************************************/
 
 IndexShardsIVF::IndexShardsIVF(
-        Index* quantizer,
-        size_t nlist,
+        Index* quantizer_,
+        size_t nlist_,
         bool threaded,
-        bool successive_ids)
-        : IndexShardsTemplate<Index>(quantizer->d, threaded, successive_ids),
-          Level1Quantizer(quantizer, nlist) {
-    is_trained = quantizer->is_trained && quantizer->ntotal == nlist;
+        bool successive_ids_)
+        : IndexShardsTemplate<Index>(quantizer_->d, threaded, successive_ids_),
+          Level1Quantizer(quantizer_, nlist_) {
+    is_trained = quantizer_->is_trained &&
+            quantizer_->ntotal == static_cast<idx_t>(nlist_);
 }
 
 void IndexShardsIVF::addIndex(Index* index) {
@@ -73,11 +74,11 @@ void IndexShardsIVF::train(idx_t n, const component_t* x) {
     for (size_t i = 0; i < indices_.size(); i++) {
         Index* index = indices_[i].first;
         auto index_ivf = dynamic_cast<IndexIVFInterface*>(index);
-        Index* quantizer = index_ivf->quantizer;
-        if (!quantizer->is_trained) {
-            quantizer->train(nlist, centroids.data());
+        Index* sub_quantizer = index_ivf->quantizer;
+        if (!sub_quantizer->is_trained) {
+            sub_quantizer->train(nlist, centroids.data());
         }
-        quantizer->add(nlist, centroids.data());
+        sub_quantizer->add(nlist, centroids.data());
         // finish training
         index->train(n, x);
     }
@@ -132,9 +133,9 @@ void IndexShardsIVF::add_with_ids(
         }
         ids = aids.data();
     }
-    idx_t d = this->d;
+    idx_t cur_d = this->d;
 
-    auto fn = [n, ids, x, nshard, d, Iq](int no, Index* index) {
+    auto fn = [n, ids, x, nshard, cur_d, Iq](int no, Index* index) {
         idx_t i0 = (idx_t)no * n / nshard;
         idx_t i1 = ((idx_t)no + 1) * n / nshard;
         auto index_ivf = dynamic_cast<IndexIVF*>(index);
@@ -144,7 +145,10 @@ void IndexShardsIVF::add_with_ids(
         }
 
         index_ivf->add_core(
-                i1 - i0, x + i0 * d, ids ? ids + i0 : nullptr, Iq.data() + i0);
+                i1 - i0,
+                x + i0 * cur_d,
+                ids ? ids + i0 : nullptr,
+                Iq.data() + i0);
 
         if (index->verbose) {
             printf("end add shard %d on %" PRId64 " points\n", no, i1 - i0);
@@ -179,7 +183,7 @@ void IndexShardsIVF::search(
 
     quantizer->search(n, x, nprobe, Dq.data(), Iq.data());
 
-    int64_t nshard = this->count();
+    int nshard = this->count();
 
     std::vector<distance_t> all_distances(nshard * k * n);
     std::vector<idx_t> all_labels(nshard * k * n);
@@ -199,7 +203,9 @@ void IndexShardsIVF::search(
 
         auto index = dynamic_cast<const IndexIVFInterface*>(indexIn);
 
-        FAISS_THROW_IF_NOT_MSG(index->nprobe == nprobe, "inconsistent nprobe");
+        FAISS_THROW_IF_NOT_MSG(
+                index->nprobe == static_cast<size_t>(nprobe),
+                "inconsistent nprobe");
 
         index->search_preassigned(
                 n,

data/vendor/faiss/faiss/MatrixStats.cpp

@@ -52,7 +52,7 @@ void MatrixStats::PerDimStats::compute_mean_std() {
     if (var < 0) {
         var = 0;
     }
-    stddev = sqrt(var);
+    stddev = std::sqrt(var);
 }
 
 void MatrixStats::do_comment(const char* fmt, ...) {
@@ -67,7 +67,8 @@ void MatrixStats::do_comment(const char* fmt, ...) {
     buf += size;
 }
 
-MatrixStats::MatrixStats(size_t n, size_t d, const float* x) : n(n), d(d) {
+MatrixStats::MatrixStats(size_t n_in, size_t d_in, const float* x)
+        : n(n_in), d(d_in) {
     std::vector<char> comment_buf(10000);
     buf = comment_buf.data();
     nbuf = comment_buf.size();
@@ -162,8 +163,8 @@ MatrixStats::MatrixStats(size_t n, size_t d, const float* x) : n(n), d(d) {
     }
 
     { // norm stats
-        min_norm2 = sqrt(min_norm2);
-        max_norm2 = sqrt(max_norm2);
+        min_norm2 = std::sqrt(min_norm2);
+        max_norm2 = std::sqrt(max_norm2);
         do_comment(
                 "range of L2 norms=[%g, %g] (%zd null vectors)\n",
                 min_norm2,

data/vendor/faiss/faiss/MetaIndexes.cpp

@@ -27,8 +27,8 @@ namespace faiss {
  * IndexSplitVectors implementation
  *******************************************************/
 
-IndexSplitVectors::IndexSplitVectors(idx_t d, bool threaded)
-        : Index(d), own_fields(false), threaded(threaded), sum_d(0) {}
+IndexSplitVectors::IndexSplitVectors(idx_t d_in, bool threaded_in)
+        : Index(d_in), own_fields(false), threaded(threaded_in), sum_d(0) {}
 
 void IndexSplitVectors::add_sub_index(Index* index) {
     sub_indexes.push_back(index);
@@ -44,7 +44,7 @@ void IndexSplitVectors::sync_with_sub_indexes() {
     metric_type = index0->metric_type;
     is_trained = index0->is_trained;
     ntotal = index0->ntotal;
-    for (int i = 1; i < sub_indexes.size(); i++) {
+    for (size_t i = 1; i < sub_indexes.size(); i++) {
         Index* index = sub_indexes[i];
         FAISS_THROW_IF_NOT(metric_type == index->metric_type);
         FAISS_THROW_IF_NOT(ntotal == index->ntotal);
@@ -88,7 +88,7 @@ void IndexSplitVectors::search(
                            n);
                 }
                 const Index* sub_index = index->sub_indexes[no];
-                int64_t sub_d = sub_index->d, d = index->d;
+                int64_t sub_d = sub_index->d;
                 idx_t ofs = 0;
                 for (int i = 0; i < no; i++) {
                     ofs += index->sub_indexes[i]->d;
@@ -107,17 +107,18 @@ void IndexSplitVectors::search(
             };
 
     if (!threaded) {
-        for (int i = 0; i < nshard; i++) {
-            query_func(i);
+        for (int64_t i = 0; i < nshard; i++) {
+            query_func(static_cast<int>(i));
         }
     } else {
         std::vector<std::unique_ptr<WorkerThread>> threads;
         std::vector<std::future<bool>> v;
 
-        for (int i = 0; i < nshard; i++) {
+        for (int64_t i = 0; i < nshard; i++) {
             threads.emplace_back(new WorkerThread());
             WorkerThread* wt = threads.back().get();
-            v.emplace_back(wt->add([i, query_func]() { query_func(i); }));
+            v.emplace_back(wt->add(
+                    [i, query_func]() { query_func(static_cast<int>(i)); }));
         }
 
         // Blocking wait for completion
@@ -155,7 +156,7 @@ void IndexSplitVectors::reset() {
 
 IndexSplitVectors::~IndexSplitVectors() {
     if (own_fields) {
-        for (int s = 0; s < sub_indexes.size(); s++) {
+        for (size_t s = 0; s < sub_indexes.size(); s++) {
             delete sub_indexes[s];
         }
     }
@@ -166,12 +167,12 @@ IndexSplitVectors::~IndexSplitVectors() {
  */
 
 IndexRandom::IndexRandom(
-        idx_t d,
-        idx_t ntotal,
-        int64_t seed,
-        MetricType metric_type)
-        : Index(d, metric_type), seed(seed) {
-    this->ntotal = ntotal;
+        idx_t d_in,
+        idx_t ntotal_in,
+        int64_t seed_in,
+        MetricType metric_type_in)
+        : Index(d_in, metric_type_in), seed(seed_in) {
+    this->ntotal = ntotal_in;
     is_trained = true;
 }
 
@@ -216,7 +217,8 @@ void IndexRandom::search(
                 perm[j] = j;
             }
             for (int j = 0; j < k; j++) {
-                std::swap(perm[j], perm[rng.rand_int(ntotal)]);
+                std::swap(
+                        perm[j], perm[rng.rand_int(static_cast<int>(ntotal))]);
                 I[j] = perm[j];
             }
         }
@@ -234,7 +236,7 @@ void IndexRandom::search(
 
 void IndexRandom::reconstruct(idx_t key, float* recons) const {
     RandomGenerator rng(seed + 123332 + key);
-    for (size_t i = 0; i < d; i++) {
+    for (int i = 0; i < d; i++) {
         recons[i] = rng.rand_float();
     }
 }

data/vendor/faiss/faiss/MetaIndexes.h

@@ -22,7 +22,7 @@ namespace faiss {
  * used to distribute a MultiIndexQuantizer
  */
 struct IndexSplitVectors : Index {
-    bool own_fields;
+    bool own_fields = false;
     bool threaded;
     std::vector<Index*> sub_indexes;
     idx_t sum_d; /// sum of dimensions seen so far

data/vendor/faiss/faiss/MetricType.h

@@ -24,14 +24,15 @@ namespace faiss {
 /// (brute-force) indices supporting additional metric types for vector
 /// comparison.
 ///
-/// NOTE: when adding or removing values, update metric_type_from_int() below.
+/// NOTE: when adding or removing values, update metric_type_from_int()
+///       and metric_type_count() below.
 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 a faiss::Index
-                              /// metric_arg
+    METRIC_INNER_PRODUCT, ///< maximum inner product search
+    METRIC_L2,            ///< squared L2 search
+    METRIC_L1,            ///< L1 (aka cityblock)
+    METRIC_Linf,          ///< infinity distance
+    METRIC_Lp,            ///< L_p distance, p is given by a faiss::Index
+                          /// metric_arg
 
     /// some additional metrics defined in scipy.spatial.distance
     METRIC_Canberra = 20,
@@ -68,6 +69,12 @@ inline MetricType metric_type_from_int(int x) {
     return static_cast<MetricType>(x);
 }
 
+/// Count of entries in the MetricType enum.
+constexpr size_t metric_type_count() {
+    return (METRIC_Lp - METRIC_INNER_PRODUCT) + 1 +
+            (METRIC_GOWER - METRIC_Canberra) + 1;
+}
+
 } // namespace faiss
 
 #endif