faiss 0.2.4 → 0.2.5

data/vendor/faiss/faiss/IndexAdditiveQuantizerFastScan.h

@@ -0,0 +1,199 @@
+/**
+ * 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/IndexAdditiveQuantizer.h>
+#include <faiss/IndexFastScan.h>
+#include <faiss/impl/AdditiveQuantizer.h>
+#include <faiss/impl/ProductAdditiveQuantizer.h>
+#include <faiss/utils/AlignedTable.h>
+
+namespace faiss {
+
+/** Fast scan version of IndexAQ. Works for 4-bit 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.
+ *
+ * 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 IndexAdditiveQuantizerFastScan : IndexFastScan {
+    AdditiveQuantizer* aq;
+    using Search_type_t = AdditiveQuantizer::Search_type_t;
+
+    bool rescale_norm = true;
+    int norm_scale = 1;
+
+    // max number of training vectors
+    size_t max_train_points = 0;
+
+    explicit IndexAdditiveQuantizerFastScan(
+            AdditiveQuantizer* aq,
+            MetricType metric = METRIC_L2,
+            int bbs = 32);
+
+    void init(
+            AdditiveQuantizer* aq,
+            MetricType metric = METRIC_L2,
+            int bbs = 32);
+
+    IndexAdditiveQuantizerFastScan();
+
+    ~IndexAdditiveQuantizerFastScan() override;
+
+    /// build from an existing IndexAQ
+    explicit IndexAdditiveQuantizerFastScan(
+            const IndexAdditiveQuantizer& orig,
+            int bbs = 32);
+
+    void train(idx_t n, const float* x) override;
+
+    void estimate_norm_scale(idx_t n, const float* x);
+
+    void compute_codes(uint8_t* codes, idx_t n, const float* x) const override;
+
+    void compute_float_LUT(float* lut, idx_t n, const float* x) const override;
+
+    void search(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
+
+    /** Decode a set of vectors.
+     *
+     *  NOTE: The codes in the IndexAdditiveQuantizerFastScan object are non-
+     *        contiguous. But this method requires a contiguous representation.
+     *
+     * @param n       number of vectors
+     * @param bytes   input encoded vectors, size n * code_size
+     * @param x       output vectors, size n * d
+     */
+    void sa_decode(idx_t n, const uint8_t* bytes, float* x) const override;
+};
+
+/** Index based on a residual quantizer. Stored vectors are
+ * approximated by residual quantization codes.
+ * Can also be used as a codec
+ */
+struct IndexResidualQuantizerFastScan : IndexAdditiveQuantizerFastScan {
+    /// The residual quantizer used to encode the vectors
+    ResidualQuantizer rq;
+
+    /** Constructor.
+     *
+     * @param d      dimensionality of the input vectors
+     * @param M      number of subquantizers
+     * @param nbits  number of bit per subvector index
+     * @param metric  metric type
+     * @param search_type AQ search type
+     */
+    IndexResidualQuantizerFastScan(
+            int d,        ///< dimensionality of the input vectors
+            size_t M,     ///< number of subquantizers
+            size_t nbits, ///< number of bit per subvector index
+            MetricType metric = METRIC_L2,
+            Search_type_t search_type = AdditiveQuantizer::ST_norm_rq2x4,
+            int bbs = 32);
+
+    IndexResidualQuantizerFastScan();
+};
+
+/** Index based on a local search quantizer. Stored vectors are
+ * approximated by local search quantization codes.
+ * Can also be used as a codec
+ */
+struct IndexLocalSearchQuantizerFastScan : IndexAdditiveQuantizerFastScan {
+    LocalSearchQuantizer lsq;
+
+    /** Constructor.
+     *
+     * @param d      dimensionality of the input vectors
+     * @param M      number of subquantizers
+     * @param nbits  number of bit per subvector index
+     * @param metric  metric type
+     * @param search_type AQ search type
+     */
+    IndexLocalSearchQuantizerFastScan(
+            int d,        ///< dimensionality of the input vectors
+            size_t M,     ///< number of subquantizers
+            size_t nbits, ///< number of bit per subvector index
+            MetricType metric = METRIC_L2,
+            Search_type_t search_type = AdditiveQuantizer::ST_norm_lsq2x4,
+            int bbs = 32);
+
+    IndexLocalSearchQuantizerFastScan();
+};
+
+/** Index based on a product residual quantizer. Stored vectors are
+ * approximated by product residual quantization codes.
+ * Can also be used as a codec
+ */
+struct IndexProductResidualQuantizerFastScan : IndexAdditiveQuantizerFastScan {
+    /// The product residual quantizer used to encode the vectors
+    ProductResidualQuantizer prq;
+
+    /** Constructor.
+     *
+     * @param d      dimensionality of the input vectors
+     * @param nsplits  number of residual quantizers
+     * @param Msub     number of subquantizers per RQ
+     * @param nbits  number of bit per subvector index
+     * @param metric  metric type
+     * @param search_type AQ search type
+     */
+    IndexProductResidualQuantizerFastScan(
+            int d,          ///< dimensionality of the input vectors
+            size_t nsplits, ///< number of residual quantizers
+            size_t Msub,    ///< number of subquantizers per RQ
+            size_t nbits,   ///< number of bit per subvector index
+            MetricType metric = METRIC_L2,
+            Search_type_t search_type = AdditiveQuantizer::ST_norm_rq2x4,
+            int bbs = 32);
+
+    IndexProductResidualQuantizerFastScan();
+};
+
+/** Index based on a product local search quantizer. Stored vectors are
+ * approximated by product local search quantization codes.
+ * Can also be used as a codec
+ */
+struct IndexProductLocalSearchQuantizerFastScan
+        : IndexAdditiveQuantizerFastScan {
+    /// The product local search quantizer used to encode the vectors
+    ProductLocalSearchQuantizer plsq;
+
+    /** Constructor.
+     *
+     * @param d      dimensionality of the input vectors
+     * @param nsplits  number of local search quantizers
+     * @param Msub     number of subquantizers per LSQ
+     * @param nbits  number of bit per subvector index
+     * @param metric  metric type
+     * @param search_type AQ search type
+     */
+    IndexProductLocalSearchQuantizerFastScan(
+            int d,          ///< dimensionality of the input vectors
+            size_t nsplits, ///< number of local search quantizers
+            size_t Msub,    ///< number of subquantizers per LSQ
+            size_t nbits,   ///< number of bit per subvector index
+            MetricType metric = METRIC_L2,
+            Search_type_t search_type = AdditiveQuantizer::ST_norm_rq2x4,
+            int bbs = 32);
+
+    IndexProductLocalSearchQuantizerFastScan();
+};
+
+} // namespace faiss

data/vendor/faiss/faiss/IndexBinary.cpp

@@ -21,8 +21,12 @@ void IndexBinary::train(idx_t, const uint8_t*) {
     // Does nothing by default.
 }
 
-void IndexBinary::range_search(idx_t, const uint8_t*, int, RangeSearchResult*)
-        const {
+void IndexBinary::range_search(
+        idx_t,
+        const uint8_t*,
+        int,
+        RangeSearchResult*,
+        const SearchParameters*) const {
     FAISS_THROW_MSG("range search not implemented");
 }
 
@@ -57,10 +61,11 @@ void IndexBinary::search_and_reconstruct(
         idx_t k,
         int32_t* distances,
         idx_t* labels,
-        uint8_t* recons) const {
+        uint8_t* recons,
+        const SearchParameters* params) const {
     FAISS_THROW_IF_NOT(k > 0);
 
-    search(n, x, k, distances, labels);
+    search(n, x, k, distances, labels, params);
     for (idx_t i = 0; i < n; ++i) {
         for (idx_t j = 0; j < k; ++j) {
             idx_t ij = i * k + j;
@@ -82,4 +87,15 @@ void IndexBinary::display() const {
            ntotal);
 }
 
+void IndexBinary::merge_from(
+        IndexBinary& /* otherIndex */,
+        idx_t /* add_id */) {
+    FAISS_THROW_MSG("merge_from() not implemented");
+}
+
+void IndexBinary::check_compatible_for_merge(
+        const IndexBinary& /* otherIndex */) const {
+    FAISS_THROW_MSG("check_compatible_for_merge() not implemented");
+}
+
 } // namespace faiss

data/vendor/faiss/faiss/IndexBinary.h

@@ -97,7 +97,8 @@ struct IndexBinary {
             const uint8_t* x,
             idx_t k,
             int32_t* distances,
-            idx_t* labels) const = 0;
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const = 0;
 
     /** Query n vectors of dimension d to the index.
      *
@@ -117,7 +118,8 @@ struct IndexBinary {
             idx_t n,
             const uint8_t* x,
             int radius,
-            RangeSearchResult* result) const;
+            RangeSearchResult* result,
+            const SearchParameters* params = nullptr) const;
 
     /** Return the indexes of the k vectors closest to the query x.
      *
@@ -164,10 +166,23 @@ struct IndexBinary {
             idx_t k,
             int32_t* distances,
             idx_t* labels,
-            uint8_t* recons) const;
+            uint8_t* recons,
+            const SearchParameters* params = nullptr) const;
 
     /** Display the actual class name and some more info. */
     void display() const;
+
+    /** moves the entries from another dataset to self.
+     * On output, other is empty.
+     * add_id is added to all moved ids
+     * (for sequential ids, this would be this->ntotal) */
+    virtual void merge_from(IndexBinary& otherIndex, idx_t add_id = 0);
+
+    /** check that the two indexes are compatible (ie, they are
+     * trained in the same way and have the same
+     * parameters). Otherwise throw. */
+    virtual void check_compatible_for_merge(
+            const IndexBinary& otherIndex) const;
 };
 
 } // namespace faiss

data/vendor/faiss/faiss/IndexBinaryFlat.cpp

@@ -11,6 +11,7 @@
 
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/IDSelector.h>
 #include <faiss/utils/Heap.h>
 #include <faiss/utils/hamming.h>
 #include <faiss/utils/utils.h>
@@ -35,7 +36,10 @@ void IndexBinaryFlat::search(
         const uint8_t* x,
         idx_t k,
         int32_t* 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);
 
     const idx_t block_size = query_batch_size;
@@ -101,7 +105,10 @@ void IndexBinaryFlat::range_search(
         idx_t n,
         const uint8_t* x,
         int radius,
-        RangeSearchResult* result) const {
+        RangeSearchResult* result,
+        const SearchParameters* params) const {
+    FAISS_THROW_IF_NOT_MSG(
+            !params, "search params not supported for this index");
     hamming_range_search(x, xb.data(), n, ntotal, radius, code_size, result);
 }
 

data/vendor/faiss/faiss/IndexBinaryFlat.h

@@ -39,13 +39,15 @@ struct IndexBinaryFlat : IndexBinary {
             const uint8_t* x,
             idx_t k,
             int32_t* distances,
-            idx_t* labels) const override;
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
 
     void range_search(
             idx_t n,
             const uint8_t* x,
             int radius,
-            RangeSearchResult* result) const override;
+            RangeSearchResult* result,
+            const SearchParameters* params = nullptr) const override;
 
     void reconstruct(idx_t key, uint8_t* recons) const override;
 

data/vendor/faiss/faiss/IndexBinaryFromFloat.cpp

@@ -52,7 +52,10 @@ void IndexBinaryFromFloat::search(
         const uint8_t* x,
         idx_t k,
         int32_t* 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);
 
     constexpr idx_t bs = 32768;

data/vendor/faiss/faiss/IndexBinaryFromFloat.h

@@ -43,7 +43,8 @@ struct IndexBinaryFromFloat : IndexBinary {
             const uint8_t* x,
             idx_t k,
             int32_t* distances,
-            idx_t* labels) const override;
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
 
     void train(idx_t n, const uint8_t* x) override;
 };

data/vendor/faiss/faiss/IndexBinaryHNSW.cpp

@@ -26,6 +26,7 @@
 
 #include <faiss/IndexBinaryFlat.h>
 #include <faiss/impl/AuxIndexStructures.h>
+#include <faiss/impl/DistanceComputer.h>
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/utils/Heap.h>
 #include <faiss/utils/hamming.h>
@@ -194,7 +195,10 @@ void IndexBinaryHNSW::search(
         const uint8_t* x,
         idx_t k,
         int32_t* 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);
 
 #pragma omp parallel

data/vendor/faiss/faiss/IndexBinaryHNSW.h

@@ -47,7 +47,8 @@ struct IndexBinaryHNSW : IndexBinary {
             const uint8_t* x,
             idx_t k,
             int32_t* distances,
-            idx_t* labels) const override;
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
 
     void reconstruct(idx_t key, uint8_t* recons) const override;
 

data/vendor/faiss/faiss/IndexBinaryHash.cpp

@@ -12,6 +12,7 @@
 #include <cinttypes>
 #include <cstdio>
 #include <memory>
+#include <unordered_set>
 
 #include <faiss/utils/hamming.h>
 #include <faiss/utils/utils.h>
@@ -216,7 +217,10 @@ void IndexBinaryHash::range_search(
         idx_t n,
         const uint8_t* x,
         int radius,
-        RangeSearchResult* result) const {
+        RangeSearchResult* result,
+        const SearchParameters* params) const {
+    FAISS_THROW_IF_NOT_MSG(
+            !params, "search params not supported for this index");
     size_t nlist = 0, ndis = 0, n0 = 0;
 
 #pragma omp parallel if (n > 100) reduction(+ : ndis, n0, nlist)
@@ -244,7 +248,10 @@ void IndexBinaryHash::search(
         const uint8_t* x,
         idx_t k,
         int32_t* 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);
 
     using HeapForL2 = CMax<int32_t, idx_t>;
@@ -431,7 +438,10 @@ void IndexBinaryMultiHash::range_search(
         idx_t n,
         const uint8_t* x,
         int radius,
-        RangeSearchResult* result) const {
+        RangeSearchResult* result,
+        const SearchParameters* params) const {
+    FAISS_THROW_IF_NOT_MSG(
+            !params, "search params not supported for this index");
     size_t nlist = 0, ndis = 0, n0 = 0;
 
 #pragma omp parallel if (n > 100) reduction(+ : ndis, n0, nlist)
@@ -459,7 +469,10 @@ void IndexBinaryMultiHash::search(
         const uint8_t* x,
         idx_t k,
         int32_t* 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);
 
     using HeapForL2 = CMax<int32_t, idx_t>;

data/vendor/faiss/faiss/IndexBinaryHash.h

@@ -50,14 +50,16 @@ struct IndexBinaryHash : IndexBinary {
             idx_t n,
             const uint8_t* x,
             int radius,
-            RangeSearchResult* result) const override;
+            RangeSearchResult* result,
+            const SearchParameters* params = nullptr) const override;
 
     void search(
             idx_t n,
             const uint8_t* x,
             idx_t k,
             int32_t* distances,
-            idx_t* labels) const override;
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
 
     void display() const;
     size_t hashtable_size() const;
@@ -107,14 +109,16 @@ struct IndexBinaryMultiHash : IndexBinary {
             idx_t n,
             const uint8_t* x,
             int radius,
-            RangeSearchResult* result) const override;
+            RangeSearchResult* result,
+            const SearchParameters* params = nullptr) const override;
 
     void search(
             idx_t n,
             const uint8_t* x,
             idx_t k,
             int32_t* distances,
-            idx_t* labels) const override;
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
 
     size_t hashtable_size() const;
 };

data/vendor/faiss/faiss/IndexBinaryIVF.cpp

@@ -125,7 +125,10 @@ void IndexBinaryIVF::search(
         const uint8_t* x,
         idx_t k,
         int32_t* 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(nprobe > 0);
 
@@ -175,7 +178,10 @@ void IndexBinaryIVF::search_and_reconstruct(
         idx_t k,
         int32_t* distances,
         idx_t* labels,
-        uint8_t* recons) const {
+        uint8_t* recons,
+        const SearchParameters* params) const {
+    FAISS_THROW_IF_NOT_MSG(
+            !params, "search params not supported for this index");
     const size_t nprobe = std::min(nlist, this->nprobe);
     FAISS_THROW_IF_NOT(k > 0);
     FAISS_THROW_IF_NOT(nprobe > 0);
@@ -279,22 +285,28 @@ void IndexBinaryIVF::train(idx_t n, const uint8_t* x) {
     is_trained = true;
 }
 
-void IndexBinaryIVF::merge_from(IndexBinaryIVF& other, idx_t add_id) {
-    // minimal sanity checks
-    FAISS_THROW_IF_NOT(other.d == d);
-    FAISS_THROW_IF_NOT(other.nlist == nlist);
-    FAISS_THROW_IF_NOT(other.code_size == code_size);
+void IndexBinaryIVF::check_compatible_for_merge(
+        const IndexBinary& otherIndex) const {
+    auto other = dynamic_cast<const IndexBinaryIVF*>(&otherIndex);
+    FAISS_THROW_IF_NOT(other);
+    FAISS_THROW_IF_NOT(other->d == d);
+    FAISS_THROW_IF_NOT(other->nlist == nlist);
+    FAISS_THROW_IF_NOT(other->code_size == code_size);
     FAISS_THROW_IF_NOT_MSG(
-            direct_map.no() && other.direct_map.no(),
+            direct_map.no() && other->direct_map.no(),
             "direct map copy not implemented");
     FAISS_THROW_IF_NOT_MSG(
             typeid(*this) == typeid(other),
             "can only merge indexes of the same type");
+}
 
-    invlists->merge_from(other.invlists, add_id);
-
-    ntotal += other.ntotal;
-    other.ntotal = 0;
+void IndexBinaryIVF::merge_from(IndexBinary& otherIndex, idx_t add_id) {
+    // minimal sanity checks
+    check_compatible_for_merge(otherIndex);
+    auto other = static_cast<IndexBinaryIVF*>(&otherIndex);
+    invlists->merge_from(other->invlists, add_id);
+    ntotal += other->ntotal;
+    other->ntotal = 0;
 }
 
 void IndexBinaryIVF::replace_invlists(InvertedLists* il, bool own) {
@@ -650,7 +662,10 @@ void IndexBinaryIVF::range_search(
         idx_t n,
         const uint8_t* x,
         int radius,
-        RangeSearchResult* res) const {
+        RangeSearchResult* res,
+        const SearchParameters* params) const {
+    FAISS_THROW_IF_NOT_MSG(
+            !params, "search params not supported for this index");
     const size_t nprobe = std::min(nlist, this->nprobe);
     std::unique_ptr<idx_t[]> idx(new idx_t[n * nprobe]);
     std::unique_ptr<int32_t[]> coarse_dis(new int32_t[n * nprobe]);

data/vendor/faiss/faiss/IndexBinaryIVF.h

@@ -123,13 +123,15 @@ struct IndexBinaryIVF : IndexBinary {
             const uint8_t* x,
             idx_t k,
             int32_t* distances,
-            idx_t* labels) const override;
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
 
     void range_search(
             idx_t n,
             const uint8_t* x,
             int radius,
-            RangeSearchResult* result) const override;
+            RangeSearchResult* result,
+            const SearchParameters* params = nullptr) const override;
 
     void range_search_preassigned(
             idx_t n,
@@ -167,7 +169,8 @@ struct IndexBinaryIVF : IndexBinary {
             idx_t k,
             int32_t* distances,
             idx_t* labels,
-            uint8_t* recons) const override;
+            uint8_t* recons,
+            const SearchParameters* params = nullptr) const override;
 
     /** Reconstruct a vector given the location in terms of (inv list index +
      * inv list offset) instead of the id.
@@ -184,10 +187,10 @@ struct IndexBinaryIVF : IndexBinary {
     /// Dataset manipulation functions
     size_t remove_ids(const IDSelector& sel) override;
 
-    /** moves the entries from another dataset to self. On output,
-     * other is empty. add_id is added to all moved ids (for
-     * sequential ids, this would be this->ntotal */
-    virtual void merge_from(IndexBinaryIVF& other, idx_t add_id);
+    void merge_from(IndexBinary& other, idx_t add_id) override;
+
+    void check_compatible_for_merge(
+            const IndexBinary& otherIndex) const override;
 
     size_t get_list_size(size_t list_no) const {
         return invlists->list_size(list_no);