faiss 0.2.3 → 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

@@ -19,38 +19,31 @@
 
 namespace faiss {
 
-IndexFlat::IndexFlat(idx_t d, MetricType metric) : Index(d, metric) {}
-
-void IndexFlat::add(idx_t n, const float* x) {
-    xb.insert(xb.end(), x, x + n * d);
-    ntotal += n;
-}
-
-void IndexFlat::reset() {
-    xb.clear();
-    ntotal = 0;
-}
+IndexFlat::IndexFlat(idx_t d, MetricType metric)
+        : IndexFlatCodes(sizeof(float) * d, d, metric) {}
 
 void IndexFlat::search(
         idx_t n,
         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, xb.data(), 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, xb.data(), 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, xb.data(), d, n, ntotal, metric_type, metric_arg, &res);
+                x, get_xb(), d, n, ntotal, metric_type, metric_arg, &res);
     }
 }
 
@@ -58,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, xb.data(), d, n, ntotal, radius, result);
+                    x, get_xb(), d, n, ntotal, radius, result, sel);
             break;
         case METRIC_L2:
-            range_search_L2sqr(x, xb.data(), 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");
@@ -80,49 +76,28 @@ void IndexFlat::compute_distance_subset(
         const idx_t* labels) const {
     switch (metric_type) {
         case METRIC_INNER_PRODUCT:
-            fvec_inner_products_by_idx(
-                    distances, x, xb.data(), labels, d, n, k);
+            fvec_inner_products_by_idx(distances, x, get_xb(), labels, d, n, k);
             break;
         case METRIC_L2:
-            fvec_L2sqr_by_idx(distances, x, xb.data(), labels, d, n, k);
+            fvec_L2sqr_by_idx(distances, x, get_xb(), labels, d, n, k);
             break;
         default:
             FAISS_THROW_MSG("metric type not supported");
     }
 }
 
-size_t IndexFlat::remove_ids(const IDSelector& sel) {
-    idx_t j = 0;
-    for (idx_t i = 0; i < ntotal; i++) {
-        if (sel.is_member(i)) {
-            // should be removed
-        } else {
-            if (i > j) {
-                memmove(&xb[d * j], &xb[d * i], sizeof(xb[0]) * d);
-            }
-            j++;
-        }
-    }
-    size_t nremove = ntotal - j;
-    if (nremove > 0) {
-        ntotal = j;
-        xb.resize(ntotal * d);
-    }
-    return nremove;
-}
-
 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 {
@@ -130,10 +105,13 @@ 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.xb.data()),
+              b(storage.get_xb()),
               ndis(0) {}
 
     void set_query(const float* x) override {
@@ -141,27 +119,30 @@ 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.xb.data()),
+              b(storage.get_xb()),
               ndis(0) {}
 
     void set_query(const float* x) override {
@@ -171,32 +152,31 @@ 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) {
         return new FlatIPDis(*this);
     } else {
         return get_extra_distance_computer(
-                d, metric_type, metric_arg, ntotal, xb.data());
+                d, metric_type, metric_arg, ntotal, get_xb());
     }
 }
 
 void IndexFlat::reconstruct(idx_t key, float* recons) const {
-    memcpy(recons, &(xb[key * d]), sizeof(*recons) * d);
-}
-
-/* The standalone codec interface */
-size_t IndexFlat::sa_code_size() const {
-    return sizeof(float) * d;
+    memcpy(recons, &(codes[key * code_size]), code_size);
 }
 
 void IndexFlat::sa_encode(idx_t n, const float* x, uint8_t* bytes) const {
-    memcpy(bytes, x, sizeof(float) * d * n);
+    if (n > 0) {
+        memcpy(bytes, x, sizeof(float) * d * n);
+    }
 }
 
 void IndexFlat::sa_decode(idx_t n, const uint8_t* bytes, float* x) const {
-    memcpy(x, bytes, sizeof(float) * d * n);
+    if (n > 0) {
+        memcpy(x, bytes, sizeof(float) * d * n);
+    }
 }
 
 /***************************************************
@@ -211,9 +191,9 @@ IndexFlat1D::IndexFlat1D(bool continuous_update)
 void IndexFlat1D::update_permutation() {
     perm.resize(ntotal);
     if (ntotal < 1000000) {
-        fvec_argsort(ntotal, xb.data(), (size_t*)perm.data());
+        fvec_argsort(ntotal, get_xb(), (size_t*)perm.data());
     } else {
-        fvec_argsort_parallel(ntotal, xb.data(), (size_t*)perm.data());
+        fvec_argsort_parallel(ntotal, get_xb(), (size_t*)perm.data());
     }
 }
 
@@ -233,11 +213,14 @@ 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();
 
 #pragma omp parallel for
     for (idx_t i = 0; i < n; i++) {

data/vendor/faiss/faiss/IndexFlat.h

@@ -12,33 +12,28 @@
 
 #include <vector>
 
-#include <faiss/Index.h>
+#include <faiss/IndexFlatCodes.h>
 
 namespace faiss {
 
 /** Index that stores the full vectors and performs exhaustive search */
-struct IndexFlat : Index {
-    /// database vectors, size ntotal * d
-    std::vector<float> xb;
-
+struct IndexFlat : IndexFlatCodes {
     explicit IndexFlat(idx_t d, MetricType metric = METRIC_L2);
 
-    void add(idx_t n, const float* x) override;
-
-    void reset() override;
-
     void search(
             idx_t n,
             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;
 
@@ -57,18 +52,19 @@ struct IndexFlat : Index {
             float* distances,
             const idx_t* labels) const;
 
-    /** remove some ids. NB that Because of the structure of the
-     * 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;
+    // get pointer to the floating point data
+    float* get_xb() {
+        return (float*)codes.data();
+    }
+    const float* get_xb() const {
+        return (const float*)codes.data();
+    }
 
     IndexFlat() {}
 
-    DistanceComputer* get_distance_computer() const override;
+    FlatCodesDistanceComputer* get_FlatCodesDistanceComputer() const override;
 
     /* The stanadlone codec interface (just memcopies in this case) */
-    size_t sa_code_size() const override;
-
     void sa_encode(idx_t n, const float* x, uint8_t* bytes) const override;
 
     void sa_decode(idx_t n, const uint8_t* bytes, float* x) const override;
@@ -106,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

@@ -0,0 +1,101 @@
+/**
+ * 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.
+ */
+
+#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 {
+
+IndexFlatCodes::IndexFlatCodes(size_t code_size, idx_t d, MetricType metric)
+        : Index(d, metric), code_size(code_size) {}
+
+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.data() + (ntotal * code_size));
+    ntotal += n;
+}
+
+void IndexFlatCodes::reset() {
+    codes.clear();
+    ntotal = 0;
+}
+
+size_t IndexFlatCodes::sa_code_size() const {
+    return code_size;
+}
+
+size_t IndexFlatCodes::remove_ids(const IDSelector& sel) {
+    idx_t j = 0;
+    for (idx_t i = 0; i < ntotal; i++) {
+        if (sel.is_member(i)) {
+            // should be removed
+        } else {
+            if (i > j) {
+                memmove(&codes[code_size * j],
+                        &codes[code_size * i],
+                        code_size);
+            }
+            j++;
+        }
+    }
+    size_t nremove = ntotal - j;
+    if (nremove > 0) {
+        ntotal = j;
+        codes.resize(ntotal * code_size);
+    }
+    return nremove;
+}
+
+void IndexFlatCodes::reconstruct_n(idx_t i0, idx_t ni, float* recons) const {
+    FAISS_THROW_IF_NOT(ni == 0 || (i0 >= 0 && i0 + ni <= ntotal));
+    sa_decode(ni, codes.data() + i0 * code_size, recons);
+}
+
+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

@@ -0,0 +1,59 @@
+/**
+ * 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.
+ */
+
+// -*- c++ -*-
+
+#pragma once
+
+#include <faiss/Index.h>
+#include <faiss/impl/DistanceComputer.h>
+#include <vector>
+
+namespace faiss {
+
+/** Index that encodes all vectors as fixed-size codes (size code_size). Storage
+ * is in the codes vector */
+struct IndexFlatCodes : Index {
+    size_t code_size;
+
+    /// encoded dataset, size ntotal * code_size
+    std::vector<uint8_t> codes;
+
+    IndexFlatCodes();
+
+    IndexFlatCodes(size_t code_size, idx_t d, MetricType metric = METRIC_L2);
+
+    /// default add uses sa_encode
+    void add(idx_t n, const float* x) override;
+
+    void reset() override;
+
+    /// reconstruction using the codec interface
+    void reconstruct_n(idx_t i0, idx_t ni, float* recons) const override;
+
+    void reconstruct(idx_t key, float* recons) const override;
+
+    size_t sa_code_size() const override;
+
+    /** remove some ids. NB that Because of the structure of the
+     * 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