faiss 0.5.3 → 0.6.1

data/vendor/faiss/faiss/IVFlib.cpp

@@ -18,6 +18,7 @@
 #include <faiss/MetaIndexes.h>
 #include <faiss/clone_index.h>
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/simd_dispatch.h>
 #include <faiss/index_io.h>
 #include <faiss/utils/distances.h>
 #include <faiss/utils/hamming.h>
@@ -36,7 +37,7 @@ void check_compatible_for_merge(const Index* index0, const Index* index1) {
         FAISS_THROW_IF_NOT_MSG(pt1, "both indexes should be pretransforms");
 
         FAISS_THROW_IF_NOT(pt0->chain.size() == pt1->chain.size());
-        for (int i = 0; i < pt0->chain.size(); i++) {
+        for (size_t i = 0; i < pt0->chain.size(); i++) {
             FAISS_THROW_IF_NOT(typeid(pt0->chain[i]) == typeid(pt1->chain[i]));
         }
 
@@ -124,7 +125,7 @@ void search_centroid(
         index = index_pre->index;
     }
     faiss::IndexIVF* index_ivf = dynamic_cast<faiss::IndexIVF*>(index);
-    assert(index_ivf);
+    FAISS_THROW_IF_NOT_MSG(index_ivf, "could not extract IVF index");
     index_ivf->quantizer->assign(n, x, centroid_ids);
 }
 
@@ -145,7 +146,7 @@ void search_and_return_centroids(
         index = index_pre->index;
     }
     faiss::IndexIVF* index_ivf = dynamic_cast<faiss::IndexIVF*>(index);
-    assert(index_ivf);
+    FAISS_THROW_IF_NOT_MSG(index_ivf, "could not extract IVF index");
 
     size_t nprobe = index_ivf->nprobe;
     std::vector<idx_t> cent_nos(n * nprobe);
@@ -179,9 +180,9 @@ void search_and_return_centroids(
     }
 }
 
-SlidingIndexWindow::SlidingIndexWindow(Index* index) : index(index) {
+SlidingIndexWindow::SlidingIndexWindow(Index* index_) : index(index_) {
     n_slice = 0;
-    IndexIVF* index_ivf = const_cast<IndexIVF*>(extract_index_ivf(index));
+    IndexIVF* index_ivf = const_cast<IndexIVF*>(extract_index_ivf(index_));
     ils = dynamic_cast<ArrayInvertedLists*>(index_ivf->invlists);
     FAISS_THROW_IF_NOT_MSG(
             ils, "only supports indexes with ArrayInvertedLists");
@@ -248,7 +249,7 @@ void SlidingIndexWindow::step(const Index* sub_index, bool remove_oldest) {
     IndexIVF* index_ivf = extract_index_ivf(index);
 
     if (remove_oldest && ils2) {
-        for (int i = 0; i < nlist; i++) {
+        for (size_t i = 0; i < nlist; i++) {
             std::vector<size_t>& sizesi = sizes[i];
             size_t amount_to_remove = sizesi[0];
             index_ivf->ntotal += ils2->ids[i].size() - amount_to_remove;
@@ -264,7 +265,7 @@ void SlidingIndexWindow::step(const Index* sub_index, bool remove_oldest) {
             sizesi[n_slice - 1] = ils->ids[i].size();
         }
     } else if (ils2) {
-        for (int i = 0; i < nlist; i++) {
+        for (size_t i = 0; i < nlist; i++) {
             index_ivf->ntotal += ils2->ids[i].size();
             shift_and_add(ils->ids[i], 0, ils2->ids[i]);
             shift_and_add(ils->codes[i], 0, ils2->codes[i]);
@@ -272,7 +273,7 @@ void SlidingIndexWindow::step(const Index* sub_index, bool remove_oldest) {
         }
         n_slice++;
     } else if (remove_oldest) {
-        for (int i = 0; i < nlist; i++) {
+        for (size_t i = 0; i < nlist; i++) {
             size_t amount_to_remove = sizes[i][0];
             index_ivf->ntotal -= amount_to_remove;
             remove_from_begin(ils->ids[i], amount_to_remove);
@@ -295,7 +296,8 @@ void SlidingIndexWindow::step(const Index* sub_index, bool remove_oldest) {
 ArrayInvertedLists* get_invlist_range(const Index* index, long i0, long i1) {
     const IndexIVF* ivf = extract_index_ivf(index);
 
-    FAISS_THROW_IF_NOT(0 <= i0 && i0 <= i1 && i1 <= ivf->nlist);
+    FAISS_THROW_IF_NOT(
+            0 <= i0 && i0 <= i1 && static_cast<size_t>(i1) <= ivf->nlist);
 
     const InvertedLists* src = ivf->invlists;
 
@@ -318,12 +320,14 @@ void set_invlist_range(
         ArrayInvertedLists* src) {
     IndexIVF* ivf = extract_index_ivf(index);
 
-    FAISS_THROW_IF_NOT(0 <= i0 && i0 <= i1 && i1 <= ivf->nlist);
+    FAISS_THROW_IF_NOT(
+            0 <= i0 && i0 <= i1 && static_cast<size_t>(i1) <= ivf->nlist);
 
     ArrayInvertedLists* dst = dynamic_cast<ArrayInvertedLists*>(ivf->invlists);
     FAISS_THROW_IF_NOT_MSG(dst, "only ArrayInvertedLists supported");
     FAISS_THROW_IF_NOT(
-            src->nlist == i1 - i0 && dst->code_size == src->code_size);
+            src->nlist == static_cast<size_t>(i1 - i0) &&
+            dst->code_size == src->code_size);
 
     size_t ntotal = index->ntotal;
     for (long i = i0; i < i1; i++) {
@@ -341,7 +345,7 @@ static size_t count_ndis(
         const idx_t* Iq) {
     size_t nb_dis = 0;
     const InvertedLists* il = index_ivf->invlists;
-    for (idx_t i = 0; i < n_list_scan; i++) {
+    for (size_t i = 0; i < n_list_scan; i++) {
         if (Iq[i] >= 0) {
             nb_dis += il->list_size(Iq[i]);
         }
@@ -456,12 +460,12 @@ void range_search_with_parameters(
 IndexIVFResidualQuantizer* ivf_residual_from_quantizer(
         const ResidualQuantizer& rq,
         int nlevel) {
-    FAISS_THROW_IF_NOT(nlevel > 0 && nlevel + 1 < rq.M);
+    FAISS_THROW_IF_NOT(nlevel > 0 && static_cast<size_t>(nlevel + 1) < rq.M);
 
     std::vector<size_t> nbits(nlevel);
     std::copy(rq.nbits.begin(), rq.nbits.begin() + nlevel, nbits.begin());
     std::unique_ptr<ResidualCoarseQuantizer> rcq(
-            new ResidualCoarseQuantizer(rq.d, nbits));
+            new ResidualCoarseQuantizer(static_cast<int>(rq.d), nbits));
 
     // set the coarse quantizer from the 2 first quantizers
     rcq->rq.initialize_from(rq);
@@ -476,7 +480,7 @@ IndexIVFResidualQuantizer* ivf_residual_from_quantizer(
 
     // build a IVFResidualQuantizer from that
     std::vector<size_t> nbits_refined;
-    for (int i = nlevel; i < rq.M; i++) {
+    for (size_t i = nlevel; i < rq.M; i++) {
         nbits_refined.push_back(rq.nbits[i]);
     }
     std::unique_ptr<IndexIVFResidualQuantizer> index(
@@ -488,7 +492,8 @@ IndexIVFResidualQuantizer* ivf_residual_from_quantizer(
                     faiss::METRIC_L2,
                     rq.search_type));
     index->own_fields = true;
-    rcq.release();
+    // Ownership transferred to index via own_fields = true
+    (void)rcq.release();
     index->by_residual = true;
     index->rq.initialize_from(rq, nlevel);
     index->is_trained = true;
@@ -512,39 +517,43 @@ void ivf_residual_add_from_flat_codes(
     const ResidualQuantizer& rq = index->rq;
 
     // populate inverted lists
-#pragma omp parallel if (nb > 10000)
-    {
-        std::vector<uint8_t> tmp_code(index->code_size);
-        std::vector<float> tmp(rq.d);
-        int nt = omp_get_num_threads();
-        int rank = omp_get_thread_num();
+    with_simd_level([&]<SIMDLevel SL>() {
+#pragma omp parallel
+        {
+            std::vector<uint8_t> tmp_code(index->rq.code_size);
+            std::vector<float> tmp(rq.d);
+            int nt = omp_get_num_threads();
+            int rank = omp_get_thread_num();
 
 #pragma omp for
-        for (idx_t i = 0; i < nb; i++) {
-            const uint8_t* code = &raw_codes[i * code_size];
-            BitstringReader rd(code, code_size);
-            idx_t list_no = rd.read(rcq->rq.tot_bits);
-
-            if (list_no % nt ==
-                rank) { // each thread takes care of 1/nt of the invlists
-                // copy AQ indexes one by one
-                BitstringWriter wr(tmp_code.data(), tmp_code.size());
-                for (int j = 0; j < rq.M; j++) {
-                    int nbit = rq.nbits[j];
-                    wr.write(rd.read(nbit), nbit);
+            for (idx_t i = 0; i < static_cast<idx_t>(nb); i++) {
+                const uint8_t* code = &raw_codes[i * code_size];
+                BitstringReader rd(code, code_size);
+                idx_t list_no = rd.read(static_cast<int>(rcq->rq.tot_bits));
+
+                if (list_no % nt ==
+                    rank) { // each thread takes care of 1/nt of the invlists
+                    // copy AQ indexes one by one
+                    BitstringWriter wr(tmp_code.data(), tmp_code.size());
+                    for (size_t j = 0; j < rq.M; j++) {
+                        int nbit = static_cast<int>(rq.nbits[j]);
+                        wr.write(rd.read(nbit), nbit);
+                    }
+                    // we need to recompute the norm
+                    // decode first, does not use the norm component, so that's
+                    // ok
+                    index->rq.decode(tmp_code.data(), tmp.data(), 1);
+                    float norm = fvec_norm_L2sqr<SL>(tmp.data(), rq.d);
+                    wr.write(
+                            rq.encode_norm(norm),
+                            static_cast<int>(rq.norm_bits));
+
+                    // add code to the inverted list
+                    invlists.add_entry(list_no, i, tmp_code.data());
                 }
-                // we need to recompute the norm
-                // decode first, does not use the norm component, so that's
-                // ok
-                index->rq.decode(tmp_code.data(), tmp.data(), 1);
-                float norm = fvec_norm_L2sqr(tmp.data(), rq.d);
-                wr.write(rq.encode_norm(norm), rq.norm_bits);
-
-                // add code to the inverted list
-                invlists.add_entry(list_no, i, tmp_code.data());
             }
         }
-    }
+    });
     index->ntotal += nb;
 }
 
@@ -595,7 +604,14 @@ void handle_ivf(
                     sharded_centroids[i].data());
         }
         char fname[256];
+#if defined(__GNUC__) || defined(__clang__)
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wformat-nonliteral"
+#endif
         snprintf(fname, 256, filename_template.c_str(), i);
+#if defined(__GNUC__) || defined(__clang__)
+#pragma GCC diagnostic pop
+#endif
         faiss::write_index(sharded_index, fname);
         delete sharded_index;
     }
@@ -646,7 +662,14 @@ void handle_binary_ivf(
                     sharded_centroids[i].data());
         }
         char fname[256];
+#if defined(__GNUC__) || defined(__clang__)
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wformat-nonliteral"
+#endif
         snprintf(fname, 256, filename_template.c_str(), i);
+#if defined(__GNUC__) || defined(__clang__)
+#pragma GCC diagnostic pop
+#endif
         faiss::write_index_binary(sharded_index, fname);
         delete sharded_index;
     }

data/vendor/faiss/faiss/Index.cpp

@@ -24,12 +24,20 @@ void Index::train(idx_t /*n*/, const float* /*x*/) {
     // does nothing by default
 }
 
+void Index::train_with_queries(
+        idx_t /*n*/,
+        const float* /*x*/,
+        idx_t /*n_train_q*/,
+        const float* /*xq_train*/) {
+    // does nothing by default
+}
+
 void Index::range_search(
         idx_t,
         const float*,
         float,
         RangeSearchResult*,
-        const SearchParameters* params) const {
+        const SearchParameters* /*params*/) const {
     FAISS_THROW_MSG("range search not implemented");
 }
 
@@ -47,7 +55,7 @@ void Index::add_with_ids(
 
 size_t Index::remove_ids(const IDSelector& /*sel*/) {
     FAISS_THROW_MSG("remove_ids not implemented for this type of index");
-    return -1;
+    return static_cast<size_t>(-1);
 }
 
 void Index::reconstruct(idx_t, float*) const {
@@ -105,19 +113,23 @@ void Index::search_and_reconstruct(
 }
 
 void Index::search_subset(
-        idx_t n,
-        const float* x,
-        idx_t k_base,
-        const idx_t* base_labels,
-        idx_t k,
-        float* distances,
-        idx_t* labels) const {
+        idx_t /*n*/,
+        const float* /*x*/,
+        idx_t /*k_base*/,
+        const idx_t* /*base_labels*/,
+        idx_t /*k*/,
+        float* /*distances*/,
+        idx_t* /*labels*/) const {
     FAISS_THROW_MSG("search_subset not implemented for this type of index");
 }
 
+void Index::search1(const float*, ResultHandler&, SearchParameters*) const {
+    FAISS_THROW_MSG("search1 not implemented for this type of index");
+}
+
 void Index::compute_residual(const float* x, float* residual, idx_t key) const {
     reconstruct(key, residual);
-    for (size_t i = 0; i < d; i++) {
+    for (int i = 0; i < d; i++) {
         residual[i] = x[i] - residual[i];
     }
 }
@@ -156,9 +168,10 @@ struct GenericDistanceComputer : DistanceComputer {
     size_t d;
     const Index& storage;
     std::vector<float> buf;
-    const float* q;
+    const float* q = nullptr;
 
-    explicit GenericDistanceComputer(const Index& storage) : storage(storage) {
+    explicit GenericDistanceComputer(const Index& storage_in)
+            : storage(storage_in) {
         d = storage.d;
         buf.resize(d * 2);
     }

data/vendor/faiss/faiss/Index.h

@@ -14,10 +14,9 @@
 #include <faiss/impl/FaissAssert.h>
 
 #include <cstdio>
-#include <sstream>
 
 #define FAISS_VERSION_MAJOR 1
-#define FAISS_VERSION_MINOR 13
+#define FAISS_VERSION_MINOR 14
 #define FAISS_VERSION_PATCH 2
 
 // Macro to combine the version components into a single string
@@ -55,6 +54,9 @@ namespace faiss {
 struct IDSelector;
 struct RangeSearchResult;
 struct DistanceComputer;
+template <typename T, typename TI>
+struct ResultHandlerUnordered;
+using ResultHandler = ResultHandlerUnordered<float, idx_t>;
 
 enum NumericType {
     Float32,
@@ -112,8 +114,8 @@ struct Index {
     MetricType metric_type;
     float metric_arg; ///< argument of the metric type
 
-    explicit Index(idx_t d = 0, MetricType metric = METRIC_L2)
-            : d(d),
+    explicit Index(idx_t d_in = 0, MetricType metric = METRIC_L2)
+            : d(static_cast<int>(d_in)),
               ntotal(0),
               verbose(false),
               is_trained(true),
@@ -129,6 +131,20 @@ struct Index {
      */
     virtual void train(idx_t n, const float* x);
 
+    /** Perform training on a representative set of vectors and a representative
+     * set of queries
+     *
+     * @param n         nb of training vectors
+     * @param x         training vectors, size n * d
+     * @param n_train_q nb of training queries
+     * @param xq_train  training queries, size n_train_q * d
+     */
+    virtual void train_with_queries(
+            idx_t n,
+            const float* x,
+            idx_t n_train_q,
+            const float* xq_train);
+
     virtual void train_ex(idx_t n, const void* x, NumericType numeric_type) {
         if (numeric_type == NumericType::Float32) {
             train(n, static_cast<const float*>(x));
@@ -216,6 +232,12 @@ struct Index {
         }
     }
 
+    /** search one vector with a custom result handler */
+    virtual void search1(
+            const float* x,
+            ResultHandler& handler,
+            SearchParameters* params = nullptr) const;
+
     /** query n vectors of dimension d to the index.
      *
      * return all vectors with distance < radius. Note that many

data/vendor/faiss/faiss/Index2Layer.cpp

@@ -13,10 +13,6 @@
 #include <cstdint>
 #include <cstdio>
 
-#ifdef __SSE3__
-#include <immintrin.h>
-#endif
-
 #include <algorithm>
 
 #include <faiss/IndexIVFPQ.h>
@@ -138,12 +134,13 @@ struct Distance2Level : DistanceComputer {
     size_t d;
     const Index2Layer& storage;
     std::vector<float> buf;
-    const float* q;
+    const float* q = nullptr;
 
-    const float *pq_l1_tab, *pq_l2_tab;
+    const float* pq_l1_tab = nullptr;
+    const float* pq_l2_tab = nullptr;
 
-    explicit Distance2Level(const Index2Layer& storage) : storage(storage) {
-        d = storage.d;
+    explicit Distance2Level(const Index2Layer& storage_) : storage(storage_) {
+        d = storage_.d;
         FAISS_ASSERT(storage.pq.dsub == 4);
         pq_l2_tab = storage.pq.centroids.data();
         buf.resize(2 * d);
@@ -162,46 +159,40 @@ struct Distance2Level : DistanceComputer {
 
 // well optimized for xNN+PQNN
 struct DistanceXPQ4 : Distance2Level {
-    int M, k;
+    int M = 0;
+    int k = 0;
 
-    explicit DistanceXPQ4(const Index2Layer& storage)
-            : Distance2Level(storage) {
+    explicit DistanceXPQ4(const Index2Layer& storage_)
+            : Distance2Level(storage_) {
         const IndexFlat* quantizer =
                 dynamic_cast<IndexFlat*>(storage.q1.quantizer);
 
         FAISS_ASSERT(quantizer);
-        M = storage.pq.M;
+        M = static_cast<int>(storage.pq.M);
         pq_l1_tab = quantizer->get_xb();
     }
 
     float operator()(idx_t i) override {
-#ifdef __SSE3__
         const uint8_t* code = storage.codes.data() + i * storage.code_size;
         idx_t key = 0;
         memcpy(&key, code, storage.code_size_1);
         code += storage.code_size_1;
 
-        // walking pointers
         const float* qa = q;
-        const __m128* l1_t = (const __m128*)(pq_l1_tab + d * key);
-        const __m128* pq_l2_t = (const __m128*)pq_l2_tab;
-        __m128 accu = _mm_setzero_ps();
+        const float* l1 = pq_l1_tab + d * key;
+        const float* l2 = pq_l2_tab;
+        float accu = 0;
 
         for (int m = 0; m < M; m++) {
-            __m128 qi = _mm_loadu_ps(qa);
-            __m128 recons = _mm_add_ps(l1_t[m], pq_l2_t[*code++]);
-            __m128 diff = _mm_sub_ps(qi, recons);
-            accu = _mm_add_ps(accu, _mm_mul_ps(diff, diff));
-            pq_l2_t += 256;
+            for (int j = 0; j < 4; j++) {
+                float diff = qa[j] - (l1[m * 4 + j] + l2[*code * 4 + j]);
+                accu += diff * diff;
+            }
+            code++;
+            l2 += 256 * 4;
             qa += 4;
         }
-
-        accu = _mm_hadd_ps(accu, accu);
-        accu = _mm_hadd_ps(accu, accu);
-        return _mm_cvtss_f32(accu);
-#else
-        FAISS_THROW_MSG("not implemented for non-x64 platforms");
-#endif
+        return accu;
     }
 };
 
@@ -209,15 +200,15 @@ struct DistanceXPQ4 : Distance2Level {
 struct Distance2xXPQ4 : Distance2Level {
     int M_2, mi_nbits;
 
-    explicit Distance2xXPQ4(const Index2Layer& storage)
-            : Distance2Level(storage) {
+    explicit Distance2xXPQ4(const Index2Layer& storage_)
+            : Distance2Level(storage_) {
         const MultiIndexQuantizer* mi =
                 dynamic_cast<MultiIndexQuantizer*>(storage.q1.quantizer);
 
         FAISS_ASSERT(mi);
         FAISS_ASSERT(storage.pq.M % 2 == 0);
-        M_2 = storage.pq.M / 2;
-        mi_nbits = mi->pq.nbits;
+        M_2 = static_cast<int>(storage.pq.M / 2);
+        mi_nbits = static_cast<int>(mi->pq.nbits);
         pq_l1_tab = mi->pq.centroids.data();
     }
 
@@ -226,42 +217,36 @@ struct Distance2xXPQ4 : Distance2Level {
         int64_t key01 = 0;
         memcpy(&key01, code, storage.code_size_1);
         code += storage.code_size_1;
-#ifdef __SSE3__
 
-        // walking pointers
         const float* qa = q;
-        const __m128* pq_l1_t = (const __m128*)pq_l1_tab;
-        const __m128* pq_l2_t = (const __m128*)pq_l2_tab;
-        __m128 accu = _mm_setzero_ps();
+        const float* l1 = pq_l1_tab;
+        const float* l2 = pq_l2_tab;
+        float accu = 0;
 
         for (int mi_m = 0; mi_m < 2; mi_m++) {
             int64_t l1_idx = key01 & (((int64_t)1 << mi_nbits) - 1);
-            const __m128* pq_l1 = pq_l1_t + M_2 * l1_idx;
+            const float* l1_sub = l1 + M_2 * l1_idx * 4;
 
             for (int m = 0; m < M_2; m++) {
-                __m128 qi = _mm_loadu_ps(qa);
-                __m128 recons = _mm_add_ps(pq_l1[m], pq_l2_t[*code++]);
-                __m128 diff = _mm_sub_ps(qi, recons);
-                accu = _mm_add_ps(accu, _mm_mul_ps(diff, diff));
-                pq_l2_t += 256;
+                for (int j = 0; j < 4; j++) {
+                    float diff =
+                            qa[j] - (l1_sub[m * 4 + j] + l2[*code * 4 + j]);
+                    accu += diff * diff;
+                }
+                code++;
+                l2 += 256 * 4;
                 qa += 4;
             }
-            pq_l1_t += M_2 << mi_nbits;
+            l1 += (M_2 << mi_nbits) * 4;
             key01 >>= mi_nbits;
         }
-        accu = _mm_hadd_ps(accu, accu);
-        accu = _mm_hadd_ps(accu, accu);
-        return _mm_cvtss_f32(accu);
-#else
-        FAISS_THROW_MSG("not implemented for non-x64 platforms");
-#endif
+        return accu;
     }
 };
 
 } // namespace
 
 DistanceComputer* Index2Layer::get_distance_computer() const {
-#ifdef __SSE3__
     const MultiIndexQuantizer* mi =
             dynamic_cast<MultiIndexQuantizer*>(q1.quantizer);
 
@@ -274,7 +259,6 @@ DistanceComputer* Index2Layer::get_distance_computer() const {
     if (fl && pq.dsub == 4) {
         return new DistanceXPQ4(*this);
     }
-#endif
 
     return Index::get_distance_computer();
 }