faiss 0.2.4 → 0.2.5

data/vendor/faiss/faiss/IndexIVFFastScan.h

@@ -0,0 +1,213 @@
+/**
+ * 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 <memory>
+
+#include <faiss/IndexIVF.h>
+#include <faiss/utils/AlignedTable.h>
+
+namespace faiss {
+
+/** Fast scan version of IVFPQ and IVFAQ. Works for 4-bit PQ/AQ for now.
+ *
+ * The codes in the inverted lists are not stored sequentially but
+ * grouped in blocks of size bbs. This makes it possible to very quickly
+ * compute distances with SIMD instructions.
+ *
+ * Implementations (implem):
+ * 0: auto-select implementation (default)
+ * 1: orig's search, re-implemented
+ * 2: orig's search, re-ordered by invlist
+ * 10: optimizer int16 search, collect results in heap, no qbs
+ * 11: idem, collect results in reservoir
+ * 12: optimizer int16 search, collect results in heap, uses qbs
+ * 13: idem, collect results in reservoir
+ */
+
+struct IndexIVFFastScan : IndexIVF {
+    // size of the kernel
+    int bbs; // set at build time
+
+    size_t M;
+    size_t nbits;
+    size_t ksub;
+
+    // M rounded up to a multiple of 2
+    size_t M2;
+
+    // search-time implementation
+    int implem = 0;
+    // skip some parts of the computation (for timing)
+    int skip = 0;
+    bool by_residual = false;
+
+    // batching factors at search time (0 = default)
+    int qbs = 0;
+    size_t qbs2 = 0;
+
+    IndexIVFFastScan(
+            Index* quantizer,
+            size_t d,
+            size_t nlist,
+            size_t code_size,
+            MetricType metric = METRIC_L2);
+
+    IndexIVFFastScan();
+
+    void init_fastscan(
+            size_t M,
+            size_t nbits,
+            size_t nlist,
+            MetricType metric,
+            int bbs);
+
+    ~IndexIVFFastScan() override;
+
+    /// orig's inverted lists (for debugging)
+    InvertedLists* orig_invlists = nullptr;
+
+    void add_with_ids(idx_t n, const float* x, const idx_t* xids) override;
+
+    // prepare look-up tables
+
+    virtual bool lookup_table_is_3d() const = 0;
+
+    virtual void compute_LUT(
+            size_t n,
+            const float* x,
+            const idx_t* coarse_ids,
+            const float* coarse_dis,
+            AlignedTable<float>& dis_tables,
+            AlignedTable<float>& biases) const = 0;
+
+    void compute_LUT_uint8(
+            size_t n,
+            const float* x,
+            const idx_t* coarse_ids,
+            const float* coarse_dis,
+            AlignedTable<uint8_t>& dis_tables,
+            AlignedTable<uint16_t>& biases,
+            float* normalizers) const;
+
+    void search(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
+
+    /// will just fail
+    void range_search(
+            idx_t n,
+            const float* x,
+            float radius,
+            RangeSearchResult* result,
+            const SearchParameters* params = nullptr) const override;
+
+    // internal search funcs
+
+    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 C, class Scaler>
+    void search_implem_1(
+            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_2(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels,
+            const Scaler& scaler) const;
+
+    // implem 10 and 12 are not multithreaded internally, so
+    // export search stats
+    template <class C, class Scaler>
+    void search_implem_10(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels,
+            int impl,
+            size_t* ndis_out,
+            size_t* nlist_out,
+            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,
+            size_t* ndis_out,
+            size_t* nlist_out,
+            const Scaler& scaler) const;
+
+    // implem 14 is mukltithreaded internally across nprobes and queries
+    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;
+
+    // reconstruct vectors from packed invlists
+    void reconstruct_from_offset(int64_t list_no, int64_t offset, float* recons)
+            const override;
+
+    // reconstruct orig invlists (for debugging)
+    void reconstruct_orig_invlists();
+};
+
+struct IVFFastScanStats {
+    uint64_t times[10];
+    uint64_t t_compute_distance_tables, t_round;
+    uint64_t t_copy_pack, t_scan, t_to_flat;
+    uint64_t reservoir_times[4];
+    double t_aq_encode;
+    double t_aq_norm_encode;
+
+    double Mcy_at(int i) {
+        return times[i] / (1000 * 1000.0);
+    }
+
+    double Mcy_reservoir_at(int i) {
+        return reservoir_times[i] / (1000 * 1000.0);
+    }
+    IVFFastScanStats() {
+        reset();
+    }
+    void reset() {
+        memset(this, 0, sizeof(*this));
+    }
+};
+
+FAISS_API extern IVFFastScanStats IVFFastScan_stats;
+
+} // namespace faiss

data/vendor/faiss/faiss/IndexIVFFlat.cpp

@@ -17,6 +17,8 @@
 #include <faiss/IndexFlat.h>
 
 #include <faiss/impl/AuxIndexStructures.h>
+#include <faiss/impl/IDSelector.h>
+
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/utils/distances.h>
 #include <faiss/utils/utils.h>
@@ -40,9 +42,7 @@ void IndexIVFFlat::add_core(
         idx_t n,
         const float* x,
         const int64_t* xids,
-        const int64_t* coarse_idx)
-
-{
+        const int64_t* coarse_idx) {
     FAISS_THROW_IF_NOT(is_trained);
     FAISS_THROW_IF_NOT(coarse_idx);
     assert(invlists);
@@ -118,13 +118,12 @@ void IndexIVFFlat::sa_decode(idx_t n, const uint8_t* bytes, float* x) const {
 
 namespace {
 
-template <MetricType metric, class C>
+template <MetricType metric, class C, bool use_sel>
 struct IVFFlatScanner : InvertedListScanner {
     size_t d;
 
-    IVFFlatScanner(size_t d, bool store_pairs) : d(d) {
-        this->store_pairs = store_pairs;
-    }
+    IVFFlatScanner(size_t d, bool store_pairs, const IDSelector* sel)
+            : InvertedListScanner(store_pairs, sel), d(d) {}
 
     const float* xi;
     void set_query(const float* query) override {
@@ -154,6 +153,9 @@ struct IVFFlatScanner : InvertedListScanner {
         size_t nup = 0;
         for (size_t j = 0; j < list_size; j++) {
             const float* yj = list_vecs + d * j;
+            if (use_sel && !sel->is_member(ids[j])) {
+                continue;
+            }
             float dis = metric == METRIC_INNER_PRODUCT
                     ? fvec_inner_product(xi, yj, d)
                     : fvec_L2sqr(xi, yj, d);
@@ -175,6 +177,9 @@ struct IVFFlatScanner : InvertedListScanner {
         const float* list_vecs = (const float*)codes;
         for (size_t j = 0; j < list_size; j++) {
             const float* yj = list_vecs + d * j;
+            if (use_sel && !sel->is_member(ids[j])) {
+                continue;
+            }
             float dis = metric == METRIC_INNER_PRODUCT
                     ? fvec_inner_product(xi, yj, d)
                     : fvec_L2sqr(xi, yj, d);
@@ -186,20 +191,34 @@ struct IVFFlatScanner : InvertedListScanner {
     }
 };
 
+template <bool use_sel>
+InvertedListScanner* get_InvertedListScanner1(
+        const IndexIVFFlat* ivf,
+        bool store_pairs,
+        const IDSelector* sel) {
+    if (ivf->metric_type == METRIC_INNER_PRODUCT) {
+        return new IVFFlatScanner<
+                METRIC_INNER_PRODUCT,
+                CMin<float, int64_t>,
+                use_sel>(ivf->d, store_pairs, sel);
+    } else if (ivf->metric_type == METRIC_L2) {
+        return new IVFFlatScanner<METRIC_L2, CMax<float, int64_t>, use_sel>(
+                ivf->d, store_pairs, sel);
+    } else {
+        FAISS_THROW_MSG("metric type not supported");
+    }
+}
+
 } // anonymous namespace
 
 InvertedListScanner* IndexIVFFlat::get_InvertedListScanner(
-        bool store_pairs) const {
-    if (metric_type == METRIC_INNER_PRODUCT) {
-        return new IVFFlatScanner<METRIC_INNER_PRODUCT, CMin<float, int64_t>>(
-                d, store_pairs);
-    } else if (metric_type == METRIC_L2) {
-        return new IVFFlatScanner<METRIC_L2, CMax<float, int64_t>>(
-                d, store_pairs);
+        bool store_pairs,
+        const IDSelector* sel) const {
+    if (sel) {
+        return get_InvertedListScanner1<true>(this, store_pairs, sel);
     } else {
-        FAISS_THROW_MSG("metric type not supported");
+        return get_InvertedListScanner1<false>(this, store_pairs, sel);
     }
-    return nullptr;
 }
 
 void IndexIVFFlat::reconstruct_from_offset(
@@ -447,7 +466,8 @@ void IndexIVFFlatDedup::range_search(
         idx_t,
         const float*,
         float,
-        RangeSearchResult*) const {
+        RangeSearchResult*,
+        const SearchParameters*) const {
     FAISS_THROW_MSG("not implemented");
 }
 

data/vendor/faiss/faiss/IndexIVFFlat.h

@@ -42,7 +42,8 @@ struct IndexIVFFlat : IndexIVF {
             bool include_listnos = false) const override;
 
     InvertedListScanner* get_InvertedListScanner(
-            bool store_pairs) const override;
+            bool store_pairs,
+            const IDSelector* sel) const override;
 
     void reconstruct_from_offset(int64_t list_no, int64_t offset, float* recons)
             const override;
@@ -89,7 +90,8 @@ struct IndexIVFFlatDedup : IndexIVFFlat {
             idx_t n,
             const float* x,
             float radius,
-            RangeSearchResult* result) const override;
+            RangeSearchResult* result,
+            const SearchParameters* params = nullptr) const override;
 
     /// not implemented
     void update_vectors(int nv, const idx_t* idx, const float* v) override;