faiss 0.2.4 → 0.2.5

data/vendor/faiss/faiss/IndexFastScan.h

@@ -0,0 +1,145 @@
+/**
+ * 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.
+ */
+
+#pragma once
+
+#include <faiss/Index.h>
+#include <faiss/utils/AlignedTable.h>
+
+namespace faiss {
+
+/** Fast scan version of IndexPQ and IndexAQ. Works for 4-bit PQ and AQ for now.
+ *
+ * The codes are not stored sequentially but grouped in blocks of size bbs.
+ * This makes it possible to compute distances quickly with SIMD instructions.
+ * The trailing codes (padding codes that are added to complete the last code)
+ * are garbage.
+ *
+ * Implementations:
+ * 12: blocked loop with internal loop on Q with qbs
+ * 13: same with reservoir accumulator to store results
+ * 14: no qbs with heap accumulator
+ * 15: no qbs with reservoir accumulator
+ */
+
+struct IndexFastScan : Index {
+    // implementation to select
+    int implem = 0;
+    // skip some parts of the computation (for timing)
+    int skip = 0;
+
+    // size of the kernel
+    int bbs;     // set at build time
+    int qbs = 0; // query block size 0 = use default
+
+    // vector quantizer
+    size_t M;
+    size_t nbits;
+    size_t ksub;
+    size_t code_size;
+
+    // packed version of the codes
+    size_t ntotal2;
+    size_t M2;
+
+    AlignedTable<uint8_t> codes;
+
+    // this is for testing purposes only
+    // (set when initialized by IndexPQ or IndexAQ)
+    const uint8_t* orig_codes = nullptr;
+
+    void init_fastscan(
+            int d,
+            size_t M,
+            size_t nbits,
+            MetricType metric,
+            int bbs);
+
+    IndexFastScan();
+
+    void reset() override;
+
+    void search(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
+
+    void add(idx_t n, const float* x) override;
+
+    virtual void compute_codes(uint8_t* codes, idx_t n, const float* x)
+            const = 0;
+
+    virtual void compute_float_LUT(float* lut, idx_t n, const float* x)
+            const = 0;
+
+    // called by search function
+    void compute_quantized_LUT(
+            idx_t n,
+            const float* x,
+            uint8_t* lut,
+            float* normalizers) const;
+
+    template <bool is_max, class Scaler>
+    void search_dispatch_implem(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels,
+            const Scaler& scaler) const;
+
+    template <class Cfloat, class Scaler>
+    void search_implem_234(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels,
+            const Scaler& scaler) const;
+
+    template <class C, class Scaler>
+    void search_implem_12(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels,
+            int impl,
+            const Scaler& scaler) const;
+
+    template <class C, class Scaler>
+    void search_implem_14(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels,
+            int impl,
+            const Scaler& scaler) const;
+
+    void reconstruct(idx_t key, float* recons) const override;
+    size_t remove_ids(const IDSelector& sel) override;
+    void merge_from(Index& otherIndex, idx_t add_id = 0) override;
+    void check_compatible_for_merge(const Index& otherIndex) const override;
+};
+
+struct FastScanStats {
+    uint64_t t0, t1, t2, t3;
+    FastScanStats() {
+        reset();
+    }
+    void reset() {
+        memset(this, 0, sizeof(*this));
+    }
+};
+
+FAISS_API extern FastScanStats FastScan_stats;
+
+} // namespace faiss

data/vendor/faiss/faiss/IndexFlat.cpp

@@ -27,18 +27,20 @@ void IndexFlat::search(
         const float* x,
         idx_t k,
         float* distances,
-        idx_t* labels) const {
+        idx_t* labels,
+        const SearchParameters* params) const {
+    IDSelector* sel = params ? params->sel : nullptr;
     FAISS_THROW_IF_NOT(k > 0);
 
     // we see the distances and labels as heaps
-
     if (metric_type == METRIC_INNER_PRODUCT) {
         float_minheap_array_t res = {size_t(n), size_t(k), labels, distances};
-        knn_inner_product(x, get_xb(), d, n, ntotal, &res);
+        knn_inner_product(x, get_xb(), d, n, ntotal, &res, sel);
     } else if (metric_type == METRIC_L2) {
         float_maxheap_array_t res = {size_t(n), size_t(k), labels, distances};
-        knn_L2sqr(x, get_xb(), d, n, ntotal, &res);
+        knn_L2sqr(x, get_xb(), d, n, ntotal, &res, nullptr, sel);
     } else {
+        FAISS_THROW_IF_NOT(!sel);
         float_maxheap_array_t res = {size_t(n), size_t(k), labels, distances};
         knn_extra_metrics(
                 x, get_xb(), d, n, ntotal, metric_type, metric_arg, &res);
@@ -49,14 +51,17 @@ void IndexFlat::range_search(
         idx_t n,
         const float* x,
         float radius,
-        RangeSearchResult* result) const {
+        RangeSearchResult* result,
+        const SearchParameters* params) const {
+    IDSelector* sel = params ? params->sel : nullptr;
+
     switch (metric_type) {
         case METRIC_INNER_PRODUCT:
             range_search_inner_product(
-                    x, get_xb(), d, n, ntotal, radius, result);
+                    x, get_xb(), d, n, ntotal, radius, result, sel);
             break;
         case METRIC_L2:
-            range_search_L2sqr(x, get_xb(), d, n, ntotal, radius, result);
+            range_search_L2sqr(x, get_xb(), d, n, ntotal, radius, result, sel);
             break;
         default:
             FAISS_THROW_MSG("metric type not supported");
@@ -83,16 +88,16 @@ void IndexFlat::compute_distance_subset(
 
 namespace {
 
-struct FlatL2Dis : DistanceComputer {
+struct FlatL2Dis : FlatCodesDistanceComputer {
     size_t d;
     Index::idx_t nb;
     const float* q;
     const float* b;
     size_t ndis;
 
-    float operator()(idx_t i) override {
+    float distance_to_code(const uint8_t* code) final {
         ndis++;
-        return fvec_L2sqr(q, b + i * d, d);
+        return fvec_L2sqr(q, (float*)code, d);
     }
 
     float symmetric_dis(idx_t i, idx_t j) override {
@@ -100,7 +105,10 @@ struct FlatL2Dis : DistanceComputer {
     }
 
     explicit FlatL2Dis(const IndexFlat& storage, const float* q = nullptr)
-            : d(storage.d),
+            : FlatCodesDistanceComputer(
+                      storage.codes.data(),
+                      storage.code_size),
+              d(storage.d),
               nb(storage.ntotal),
               q(q),
               b(storage.get_xb()),
@@ -111,24 +119,27 @@ struct FlatL2Dis : DistanceComputer {
     }
 };
 
-struct FlatIPDis : DistanceComputer {
+struct FlatIPDis : FlatCodesDistanceComputer {
     size_t d;
     Index::idx_t nb;
     const float* q;
     const float* b;
     size_t ndis;
 
-    float operator()(idx_t i) override {
-        ndis++;
-        return fvec_inner_product(q, b + i * d, d);
-    }
-
     float symmetric_dis(idx_t i, idx_t j) override {
         return fvec_inner_product(b + j * d, b + i * d, d);
     }
 
+    float distance_to_code(const uint8_t* code) final {
+        ndis++;
+        return fvec_inner_product(q, (float*)code, d);
+    }
+
     explicit FlatIPDis(const IndexFlat& storage, const float* q = nullptr)
-            : d(storage.d),
+            : FlatCodesDistanceComputer(
+                      storage.codes.data(),
+                      storage.code_size),
+              d(storage.d),
               nb(storage.ntotal),
               q(q),
               b(storage.get_xb()),
@@ -141,7 +152,7 @@ struct FlatIPDis : DistanceComputer {
 
 } // namespace
 
-DistanceComputer* IndexFlat::get_distance_computer() const {
+FlatCodesDistanceComputer* IndexFlat::get_FlatCodesDistanceComputer() const {
     if (metric_type == METRIC_L2) {
         return new FlatL2Dis(*this);
     } else if (metric_type == METRIC_INNER_PRODUCT) {
@@ -202,9 +213,11 @@ void IndexFlat1D::search(
         const float* x,
         idx_t k,
         float* distances,
-        idx_t* labels) const {
+        idx_t* labels,
+        const SearchParameters* params) const {
+    FAISS_THROW_IF_NOT_MSG(
+            !params, "search params not supported for this index");
     FAISS_THROW_IF_NOT(k > 0);
-
     FAISS_THROW_IF_NOT_MSG(
             perm.size() == ntotal, "Call update_permutation before search");
     const float* xb = get_xb();

data/vendor/faiss/faiss/IndexFlat.h

@@ -25,13 +25,15 @@ struct IndexFlat : IndexFlatCodes {
             const float* x,
             idx_t k,
             float* distances,
-            idx_t* labels) const override;
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
 
     void range_search(
             idx_t n,
             const float* x,
             float radius,
-            RangeSearchResult* result) const override;
+            RangeSearchResult* result,
+            const SearchParameters* params = nullptr) const override;
 
     void reconstruct(idx_t key, float* recons) const override;
 
@@ -60,7 +62,7 @@ struct IndexFlat : IndexFlatCodes {
 
     IndexFlat() {}
 
-    DistanceComputer* get_distance_computer() const override;
+    FlatCodesDistanceComputer* get_FlatCodesDistanceComputer() const override;
 
     /* The stanadlone codec interface (just memcopies in this case) */
     void sa_encode(idx_t n, const float* x, uint8_t* bytes) const override;
@@ -100,7 +102,8 @@ struct IndexFlat1D : IndexFlatL2 {
             const float* x,
             idx_t k,
             float* distances,
-            idx_t* labels) const override;
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
 };
 
 } // namespace faiss

data/vendor/faiss/faiss/IndexFlatCodes.cpp

@@ -8,7 +8,9 @@
 #include <faiss/IndexFlatCodes.h>
 
 #include <faiss/impl/AuxIndexStructures.h>
+#include <faiss/impl/DistanceComputer.h>
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/IDSelector.h>
 
 namespace faiss {
 
@@ -19,8 +21,11 @@ IndexFlatCodes::IndexFlatCodes() : code_size(0) {}
 
 void IndexFlatCodes::add(idx_t n, const float* x) {
     FAISS_THROW_IF_NOT(is_trained);
+    if (n == 0) {
+        return;
+    }
     codes.resize((ntotal + n) * code_size);
-    sa_encode(n, x, &codes[ntotal * code_size]);
+    sa_encode(n, x, codes.data() + (ntotal * code_size));
     ntotal += n;
 }
 
@@ -64,4 +69,33 @@ void IndexFlatCodes::reconstruct(idx_t key, float* recons) const {
     reconstruct_n(key, 1, recons);
 }
 
+FlatCodesDistanceComputer* IndexFlatCodes::get_FlatCodesDistanceComputer()
+        const {
+    FAISS_THROW_MSG("not implemented");
+}
+
+void IndexFlatCodes::check_compatible_for_merge(const Index& otherIndex) const {
+    // minimal sanity checks
+    const IndexFlatCodes* other =
+            dynamic_cast<const IndexFlatCodes*>(&otherIndex);
+    FAISS_THROW_IF_NOT(other);
+    FAISS_THROW_IF_NOT(other->d == d);
+    FAISS_THROW_IF_NOT(other->code_size == code_size);
+    FAISS_THROW_IF_NOT_MSG(
+            typeid(*this) == typeid(*other),
+            "can only merge indexes of the same type");
+}
+
+void IndexFlatCodes::merge_from(Index& otherIndex, idx_t add_id) {
+    FAISS_THROW_IF_NOT_MSG(add_id == 0, "cannot set ids in FlatCodes index");
+    check_compatible_for_merge(otherIndex);
+    IndexFlatCodes* other = static_cast<IndexFlatCodes*>(&otherIndex);
+    codes.resize((ntotal + other->ntotal) * code_size);
+    memcpy(codes.data() + (ntotal * code_size),
+           other->codes.data(),
+           other->ntotal * code_size);
+    ntotal += other->ntotal;
+    other->reset();
+}
+
 } // namespace faiss

data/vendor/faiss/faiss/IndexFlatCodes.h

@@ -10,6 +10,7 @@
 #pragma once
 
 #include <faiss/Index.h>
+#include <faiss/impl/DistanceComputer.h>
 #include <vector>
 
 namespace faiss {
@@ -42,6 +43,17 @@ struct IndexFlatCodes : Index {
      * indexing structure, the semantics of this operation are
      * different from the usual ones: the new ids are shifted */
     size_t remove_ids(const IDSelector& sel) override;
+
+    /** a FlatCodesDistanceComputer offers a distance_to_code method */
+    virtual FlatCodesDistanceComputer* get_FlatCodesDistanceComputer() const;
+
+    DistanceComputer* get_distance_computer() const override {
+        return get_FlatCodesDistanceComputer();
+    }
+
+    void check_compatible_for_merge(const Index& otherIndex) const override;
+
+    virtual void merge_from(Index& otherIndex, idx_t add_id = 0) override;
 };
 
 } // namespace faiss

data/vendor/faiss/faiss/IndexHNSW.cpp

@@ -202,7 +202,10 @@ void hnsw_add_vertices(
                         verbose && omp_get_thread_num() == 0 ? 0 : -1;
                 size_t counter = 0;
 
-#pragma omp for schedule(dynamic)
+                // here we should do schedule(dynamic) but this segfaults for
+                // some versions of LLVM. The performance impact should not be
+                // too large when (i1 - i0) / num_threads >> 1
+#pragma omp for schedule(static)
                 for (int i = i0; i < i1; i++) {
                     storage_idx_t pt_id = order[i];
                     dis->set_query(x + (pt_id - n0) * d);
@@ -219,7 +222,6 @@ void hnsw_add_vertices(
                         printf("  %d / %d\r", i - i0, i1 - i0);
                         fflush(stdout);
                     }
-
                     if (counter % check_period == 0) {
                         if (InterruptCallback::is_interrupted()) {
                             interrupt = true;
@@ -284,18 +286,24 @@ void IndexHNSW::search(
         const float* x,
         idx_t k,
         float* distances,
-        idx_t* labels) const
-
-{
+        idx_t* labels,
+        const SearchParameters* params_in) const {
     FAISS_THROW_IF_NOT(k > 0);
-
     FAISS_THROW_IF_NOT_MSG(
             storage,
             "Please use IndexHNSWFlat (or variants) instead of IndexHNSW directly");
+    const SearchParametersHNSW* params = nullptr;
+
+    int efSearch = hnsw.efSearch;
+    if (params_in) {
+        params = dynamic_cast<const SearchParametersHNSW*>(params_in);
+        FAISS_THROW_IF_NOT_MSG(params, "params type invalid");
+        efSearch = params->efSearch;
+    }
     size_t n1 = 0, n2 = 0, n3 = 0, ndis = 0, nreorder = 0;
 
-    idx_t check_period = InterruptCallback::get_period_hint(
-            hnsw.max_level * d * hnsw.efSearch);
+    idx_t check_period =
+            InterruptCallback::get_period_hint(hnsw.max_level * d * efSearch);
 
     for (idx_t i0 = 0; i0 < n; i0 += check_period) {
         idx_t i1 = std::min(i0 + check_period, n);
@@ -314,7 +322,7 @@ void IndexHNSW::search(
                 dis->set_query(x + i * d);
 
                 maxheap_heapify(k, simi, idxi);
-                HNSWStats stats = hnsw.search(*dis, k, idxi, simi, vt);
+                HNSWStats stats = hnsw.search(*dis, k, idxi, simi, vt, params);
                 n1 += stats.n1;
                 n2 += stats.n2;
                 n3 += stats.n3;
@@ -423,16 +431,15 @@ void IndexHNSW::search_level_0(
     FAISS_THROW_IF_NOT(nprobe > 0);
 
     storage_idx_t ntotal = hnsw.levels.size();
-    size_t n1 = 0, n2 = 0, n3 = 0, ndis = 0, nreorder = 0;
 
 #pragma omp parallel
     {
-        DistanceComputer* qdis = storage_distance_computer(storage);
-        ScopeDeleter1<DistanceComputer> del(qdis);
-
+        std::unique_ptr<DistanceComputer> qdis(
+                storage_distance_computer(storage));
+        HNSWStats search_stats;
         VisitedTable vt(ntotal);
 
-#pragma omp for reduction(+ : n1, n2, n3, ndis, nreorder)
+#pragma omp for
         for (idx_t i = 0; i < n; i++) {
             idx_t* idxi = labels + i * k;
             float* simi = distances + i * k;
@@ -440,69 +447,24 @@ void IndexHNSW::search_level_0(
             qdis->set_query(x + i * d);
             maxheap_heapify(k, simi, idxi);
 
-            if (search_type == 1) {
-                int nres = 0;
-
-                for (int j = 0; j < nprobe; j++) {
-                    storage_idx_t cj = nearest[i * nprobe + j];
-
-                    if (cj < 0)
-                        break;
-
-                    if (vt.get(cj))
-                        continue;
-
-                    int candidates_size = std::max(hnsw.efSearch, int(k));
-                    MinimaxHeap candidates(candidates_size);
-
-                    candidates.push(cj, nearest_d[i * nprobe + j]);
-
-                    HNSWStats search_stats;
-                    nres = hnsw.search_from_candidates(
-                            *qdis,
-                            k,
-                            idxi,
-                            simi,
-                            candidates,
-                            vt,
-                            search_stats,
-                            0,
-                            nres);
-                    n1 += search_stats.n1;
-                    n2 += search_stats.n2;
-                    n3 += search_stats.n3;
-                    ndis += search_stats.ndis;
-                    nreorder += search_stats.nreorder;
-                }
-            } else if (search_type == 2) {
-                int candidates_size = std::max(hnsw.efSearch, int(k));
-                candidates_size = std::max(candidates_size, nprobe);
-
-                MinimaxHeap candidates(candidates_size);
-                for (int j = 0; j < nprobe; j++) {
-                    storage_idx_t cj = nearest[i * nprobe + j];
-
-                    if (cj < 0)
-                        break;
-                    candidates.push(cj, nearest_d[i * nprobe + j]);
-                }
+            hnsw.search_level_0(
+                    *qdis.get(),
+                    k,
+                    idxi,
+                    simi,
+                    nprobe,
+                    nearest + i * nprobe,
+                    nearest_d + i * nprobe,
+                    search_type,
+                    search_stats,
+                    vt);
 
-                HNSWStats search_stats;
-                hnsw.search_from_candidates(
-                        *qdis, k, idxi, simi, candidates, vt, search_stats, 0);
-                n1 += search_stats.n1;
-                n2 += search_stats.n2;
-                n3 += search_stats.n3;
-                ndis += search_stats.ndis;
-                nreorder += search_stats.nreorder;
-            }
             vt.advance();
-
             maxheap_reorder(k, simi, idxi);
         }
+#pragma omp critical
+        { hnsw_stats.combine(search_stats); }
     }
-
-    hnsw_stats.combine({n1, n2, n3, ndis, nreorder});
 }
 
 void IndexHNSW::init_level_0_from_knngraph(
@@ -1035,8 +997,11 @@ void IndexHNSW2Level::search(
         const float* x,
         idx_t k,
         float* distances,
-        idx_t* labels) const {
+        idx_t* labels,
+        const SearchParameters* params) const {
     FAISS_THROW_IF_NOT(k > 0);
+    FAISS_THROW_IF_NOT_MSG(
+            !params, "search params not supported for this index");
 
     if (dynamic_cast<const Index2Layer*>(storage)) {
         IndexHNSW::search(n, x, k, distances, labels);
@@ -1095,74 +1060,37 @@ void IndexHNSW2Level::search(
                 }
 
                 candidates.clear();
-                // copy the upper_beam elements to candidates list
-
-                int search_policy = 2;
-
-                if (search_policy == 1) {
-                    for (int j = 0; j < hnsw.upper_beam && j < k; j++) {
-                        if (idxi[j] < 0)
-                            break;
-                        candidates.push(idxi[j], simi[j]);
-                        // search_from_candidates adds them back
-                        idxi[j] = -1;
-                        simi[j] = HUGE_VAL;
-                    }
 
-                    // reorder from sorted to heap
-                    maxheap_heapify(k, simi, idxi, simi, idxi, k);
-
-                    HNSWStats search_stats;
-                    hnsw.search_from_candidates(
-                            *dis,
-                            k,
-                            idxi,
-                            simi,
-                            candidates,
-                            vt,
-                            search_stats,
-                            0,
-                            k);
-                    n1 += search_stats.n1;
-                    n2 += search_stats.n2;
-                    n3 += search_stats.n3;
-                    ndis += search_stats.ndis;
-                    nreorder += search_stats.nreorder;
-
-                    vt.advance();
-
-                } else if (search_policy == 2) {
-                    for (int j = 0; j < hnsw.upper_beam && j < k; j++) {
-                        if (idxi[j] < 0)
-                            break;
-                        candidates.push(idxi[j], simi[j]);
-                    }
-
-                    // reorder from sorted to heap
-                    maxheap_heapify(k, simi, idxi, simi, idxi, k);
-
-                    HNSWStats search_stats;
-                    search_from_candidates_2(
-                            hnsw,
-                            *dis,
-                            k,
-                            idxi,
-                            simi,
-                            candidates,
-                            vt,
-                            search_stats,
-                            0,
-                            k);
-                    n1 += search_stats.n1;
-                    n2 += search_stats.n2;
-                    n3 += search_stats.n3;
-                    ndis += search_stats.ndis;
-                    nreorder += search_stats.nreorder;
-
-                    vt.advance();
-                    vt.advance();
+                for (int j = 0; j < hnsw.upper_beam && j < k; j++) {
+                    if (idxi[j] < 0)
+                        break;
+                    candidates.push(idxi[j], simi[j]);
                 }
 
+                // reorder from sorted to heap
+                maxheap_heapify(k, simi, idxi, simi, idxi, k);
+
+                HNSWStats search_stats;
+                search_from_candidates_2(
+                        hnsw,
+                        *dis,
+                        k,
+                        idxi,
+                        simi,
+                        candidates,
+                        vt,
+                        search_stats,
+                        0,
+                        k);
+                n1 += search_stats.n1;
+                n2 += search_stats.n2;
+                n3 += search_stats.n3;
+                ndis += search_stats.ndis;
+                nreorder += search_stats.nreorder;
+
+                vt.advance();
+                vt.advance();
+
                 maxheap_reorder(k, simi, idxi);
             }
         }

data/vendor/faiss/faiss/IndexHNSW.h

@@ -96,7 +96,8 @@ struct IndexHNSW : Index {
             const float* x,
             idx_t k,
             float* distances,
-            idx_t* labels) const override;
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
 
     void reconstruct(idx_t key, float* recons) const override;
 
@@ -180,7 +181,8 @@ struct IndexHNSW2Level : IndexHNSW {
             const float* x,
             idx_t k,
             float* distances,
-            idx_t* labels) const override;
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
 };
 
 } // namespace faiss