faiss 0.2.4 → 0.2.6

data/vendor/faiss/faiss/Index2Layer.cpp

@@ -111,7 +111,8 @@ void Index2Layer::search(
         const float* /*x*/,
         idx_t /*k*/,
         float* /*distances*/,
-        idx_t* /*labels*/) const {
+        idx_t* /*labels*/,
+        const SearchParameters* /* params */) const {
     FAISS_THROW_MSG("not implemented");
 }
 
@@ -282,10 +283,13 @@ DistanceComputer* Index2Layer::get_distance_computer() const {
 
 /* The standalone codec interface */
 
+// block size used in Index2Layer::sa_encode
+int index2layer_sa_encode_bs = 32768;
+
 void Index2Layer::sa_encode(idx_t n, const float* x, uint8_t* bytes) const {
     FAISS_THROW_IF_NOT(is_trained);
 
-    idx_t bs = 32768;
+    idx_t bs = index2layer_sa_encode_bs;
     if (n > bs) {
         for (idx_t i0 = 0; i0 < n; i0 += bs) {
             idx_t i1 = std::min(i0 + bs, n);

data/vendor/faiss/faiss/Index2Layer.h

@@ -14,6 +14,7 @@
 #include <faiss/IndexFlatCodes.h>
 #include <faiss/IndexIVF.h>
 #include <faiss/IndexPQ.h>
+#include <faiss/impl/platform_macros.h>
 
 namespace faiss {
 
@@ -56,7 +57,8 @@ struct Index2Layer : IndexFlatCodes {
             const float* x,
             idx_t k,
             float* distances,
-            idx_t* labels) const override;
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
 
     DistanceComputer* get_distance_computer() const override;
 
@@ -68,4 +70,7 @@ struct Index2Layer : IndexFlatCodes {
     void sa_decode(idx_t n, const uint8_t* bytes, float* x) const override;
 };
 
+// block size used in Index2Layer::sa_encode
+FAISS_API extern int index2layer_sa_encode_bs;
+
 } // namespace faiss

data/vendor/faiss/faiss/IndexAdditiveQuantizer.cpp

@@ -40,6 +40,90 @@ IndexAdditiveQuantizer::IndexAdditiveQuantizer(
 
 namespace {
 
+/************************************************************
+ * DistanceComputer implementation
+ ************************************************************/
+
+template <class VectorDistance>
+struct AQDistanceComputerDecompress: FlatCodesDistanceComputer {
+    std::vector<float> tmp;
+    const AdditiveQuantizer & aq;
+    VectorDistance vd;
+    size_t d;
+
+    AQDistanceComputerDecompress(const IndexAdditiveQuantizer &iaq, VectorDistance vd):
+        FlatCodesDistanceComputer(iaq.codes.data(), iaq.code_size),
+        tmp(iaq.d * 2),
+        aq(*iaq.aq),
+        vd(vd),
+        d(iaq.d)
+        {}
+
+    const float *q;
+    void set_query(const float* x) final {
+        q = x;
+    }
+
+    float symmetric_dis(idx_t i, idx_t j) final {
+        aq.decode(codes + i * d, tmp.data(), 1);
+        aq.decode(codes + j * d, tmp.data() + d, 1);
+        return vd(tmp.data(), tmp.data() + d);
+    }
+
+    float distance_to_code(const uint8_t *code) final {
+        aq.decode(code, tmp.data(), 1);
+        return vd(q, tmp.data());
+    }
+
+    virtual ~AQDistanceComputerDecompress() {}
+};
+
+
+template<bool is_IP, AdditiveQuantizer::Search_type_t st>
+struct AQDistanceComputerLUT: FlatCodesDistanceComputer {
+    std::vector<float> LUT;
+    const AdditiveQuantizer & aq;
+    size_t d;
+
+    explicit AQDistanceComputerLUT(const IndexAdditiveQuantizer &iaq):
+        FlatCodesDistanceComputer(iaq.codes.data(), iaq.code_size),
+        LUT(iaq.aq->total_codebook_size + iaq.d * 2),
+        aq(*iaq.aq),
+        d(iaq.d)
+        {}
+
+    float bias;
+    void set_query(const float* x) final {
+        // this is quite sub-optimal for multiple queries
+        aq.compute_LUT(1, x, LUT.data());
+        if (is_IP) {
+            bias = 0;
+        } else {
+            bias = fvec_norm_L2sqr(x, d);
+        }
+    }
+
+    float symmetric_dis(idx_t i, idx_t j) final {
+        float *tmp = LUT.data();
+        aq.decode(codes + i * d, tmp, 1);
+        aq.decode(codes + j * d, tmp + d, 1);
+        return fvec_L2sqr(tmp, tmp + d, d);
+    }
+
+    float distance_to_code(const uint8_t *code) final {
+        return bias + aq.compute_1_distance_LUT<is_IP, st>(code, LUT.data());
+    }
+
+    virtual ~AQDistanceComputerLUT() {}
+};
+
+
+
+/************************************************************
+ * scanning implementation for search
+ ************************************************************/
+
+
 template <class VectorDistance, class ResultHandler>
 void search_with_decompress(
         const IndexAdditiveQuantizer& ir,
@@ -111,12 +195,61 @@ void search_with_LUT(
 
 } // anonymous namespace
 
+
+FlatCodesDistanceComputer * IndexAdditiveQuantizer::get_FlatCodesDistanceComputer() const {
+
+    if (aq->search_type == AdditiveQuantizer::ST_decompress) {
+        if (metric_type == METRIC_L2) {
+            using VD = VectorDistance<METRIC_L2>;
+            VD vd = {size_t(d), metric_arg};
+            return new AQDistanceComputerDecompress<VD>(*this, vd);
+        } else if (metric_type == METRIC_INNER_PRODUCT) {
+            using VD = VectorDistance<METRIC_INNER_PRODUCT>;
+            VD vd = {size_t(d), metric_arg};
+            return new AQDistanceComputerDecompress<VD>(*this, vd);
+        } else {
+            FAISS_THROW_MSG("unsupported metric");
+        }
+    } else {
+        if (metric_type == METRIC_INNER_PRODUCT) {
+            return new AQDistanceComputerLUT<true, AdditiveQuantizer::ST_LUT_nonorm>(*this);
+        } else {
+            switch(aq->search_type) {
+#define DISPATCH(st) \
+            case AdditiveQuantizer::st: \
+                return new AQDistanceComputerLUT<false, AdditiveQuantizer::st> (*this);\
+                break;
+            DISPATCH(ST_norm_float)
+            DISPATCH(ST_LUT_nonorm)
+            DISPATCH(ST_norm_qint8)
+            DISPATCH(ST_norm_qint4)
+            DISPATCH(ST_norm_cqint4)
+            case AdditiveQuantizer::ST_norm_cqint8:
+            case AdditiveQuantizer::ST_norm_lsq2x4:
+            case AdditiveQuantizer::ST_norm_rq2x4:
+                return new AQDistanceComputerLUT<false, AdditiveQuantizer::ST_norm_cqint8> (*this);\
+                break;
+#undef DISPATCH
+            default:
+                FAISS_THROW_FMT("search type %d not supported", aq->search_type);
+            }
+        }
+    }
+}
+
+
+
+
 void IndexAdditiveQuantizer::search(
         idx_t n,
         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");
+
     if (aq->search_type == AdditiveQuantizer::ST_decompress) {
         if (metric_type == METRIC_L2) {
             using VD = VectorDistance<METRIC_L2>;
@@ -135,20 +268,23 @@ void IndexAdditiveQuantizer::search(
             search_with_LUT<true, AdditiveQuantizer::ST_LUT_nonorm> (*this, x, rh);
         } else {
             HeapResultHandler<CMax<float, idx_t> > rh(n, distances, labels, k);
-
-            if (aq->search_type == AdditiveQuantizer::ST_norm_float) {
-                search_with_LUT<false, AdditiveQuantizer::ST_norm_float> (*this, x, rh);
-            } else if (aq->search_type == AdditiveQuantizer::ST_LUT_nonorm) {
-                search_with_LUT<false, AdditiveQuantizer::ST_norm_float> (*this, x, rh);
-            } else if (aq->search_type == AdditiveQuantizer::ST_norm_qint8) {
-                search_with_LUT<false, AdditiveQuantizer::ST_norm_qint8> (*this, x, rh);
-            } else if (aq->search_type == AdditiveQuantizer::ST_norm_qint4) {
-                search_with_LUT<false, AdditiveQuantizer::ST_norm_qint4> (*this, x, rh);
-            } else if (aq->search_type == AdditiveQuantizer::ST_norm_cqint8) {
+            switch(aq->search_type) {
+#define DISPATCH(st) \
+            case AdditiveQuantizer::st: \
+                search_with_LUT<false, AdditiveQuantizer::st> (*this, x, rh);\
+                break;
+            DISPATCH(ST_norm_float)
+            DISPATCH(ST_LUT_nonorm)
+            DISPATCH(ST_norm_qint8)
+            DISPATCH(ST_norm_qint4)
+            DISPATCH(ST_norm_cqint4)
+            case AdditiveQuantizer::ST_norm_cqint8:
+            case AdditiveQuantizer::ST_norm_lsq2x4:
+            case AdditiveQuantizer::ST_norm_rq2x4:
                 search_with_LUT<false, AdditiveQuantizer::ST_norm_cqint8> (*this, x, rh);
-            } else if (aq->search_type == AdditiveQuantizer::ST_norm_cqint4) {
-                search_with_LUT<false, AdditiveQuantizer::ST_norm_cqint4> (*this, x, rh);
-            } else {
+                break;
+#undef DISPATCH
+            default:
                 FAISS_THROW_FMT("search type %d not supported", aq->search_type);
             }
         }
@@ -220,6 +356,57 @@ void IndexLocalSearchQuantizer::train(idx_t n, const float* x) {
     is_trained = true;
 }
 
+
+/**************************************************************************************
+ * IndexProductResidualQuantizer
+ **************************************************************************************/
+
+IndexProductResidualQuantizer::IndexProductResidualQuantizer(
+        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,
+        Search_type_t search_type)
+        : IndexAdditiveQuantizer(d, &prq, metric), prq(d, nsplits, Msub, nbits, search_type) {
+    code_size = prq.code_size;
+    is_trained = false;
+}
+
+IndexProductResidualQuantizer::IndexProductResidualQuantizer()
+        : IndexProductResidualQuantizer(0, 0, 0, 0) {}
+
+void IndexProductResidualQuantizer::train(idx_t n, const float* x) {
+    prq.train(n, x);
+    is_trained = true;
+}
+
+
+/**************************************************************************************
+ * IndexProductLocalSearchQuantizer
+ **************************************************************************************/
+
+IndexProductLocalSearchQuantizer::IndexProductLocalSearchQuantizer(
+        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,
+        Search_type_t search_type)
+        : IndexAdditiveQuantizer(d, &plsq, metric), plsq(d, nsplits, Msub, nbits, search_type) {
+    code_size = plsq.code_size;
+    is_trained = false;
+}
+
+IndexProductLocalSearchQuantizer::IndexProductLocalSearchQuantizer()
+        : IndexProductLocalSearchQuantizer(0, 0, 0, 0) {}
+
+void IndexProductLocalSearchQuantizer::train(idx_t n, const float* x) {
+    plsq.train(n, x);
+    is_trained = true;
+}
+
+
 /**************************************************************************************
  * AdditiveCoarseQuantizer
  **************************************************************************************/
@@ -248,6 +435,13 @@ void AdditiveCoarseQuantizer::train(idx_t n, const float* x) {
     if (verbose) {
         printf("AdditiveCoarseQuantizer::train: training on %zd vectors\n", size_t(n));
     }
+    size_t norms_size = sizeof(float) << aq->tot_bits;
+
+    FAISS_THROW_IF_NOT_MSG (
+        norms_size <= aq->max_mem_distances,
+        "the RCQ norms matrix will become too large, please reduce the number of quantization steps"
+    );
+
     aq->train(n, x);
     is_trained = true;
     ntotal = (idx_t)1 << aq->tot_bits;
@@ -268,7 +462,11 @@ void AdditiveCoarseQuantizer::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");
+
     if (metric_type == METRIC_INNER_PRODUCT) {
         aq->knn_centroids_inner_product(n, x, k, distances, labels);
     } else if (metric_type == METRIC_L2) {
@@ -321,7 +519,12 @@ void ResidualCoarseQuantizer::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");
+
     if (beam_factor < 0) {
         AdditiveCoarseQuantizer::search(n, x, k, distances, labels);
         return;

data/vendor/faiss/faiss/IndexAdditiveQuantizer.h

@@ -15,6 +15,7 @@
 
 #include <faiss/IndexFlatCodes.h>
 #include <faiss/impl/LocalSearchQuantizer.h>
+#include <faiss/impl/ProductAdditiveQuantizer.h>
 #include <faiss/impl/ResidualQuantizer.h>
 #include <faiss/impl/platform_macros.h>
 
@@ -28,8 +29,8 @@ struct IndexAdditiveQuantizer : IndexFlatCodes {
     using Search_type_t = AdditiveQuantizer::Search_type_t;
 
     explicit IndexAdditiveQuantizer(
-            idx_t d = 0,
-            AdditiveQuantizer* aq = nullptr,
+            idx_t d,
+            AdditiveQuantizer* aq,
             MetricType metric = METRIC_L2);
 
     void search(
@@ -37,12 +38,15 @@ struct IndexAdditiveQuantizer : IndexFlatCodes {
             const float* x,
             idx_t k,
             float* distances,
-            idx_t* labels) const override;
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
 
     /* The standalone codec interface */
     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;
+
+    FlatCodesDistanceComputer* get_FlatCodesDistanceComputer() const override;
 };
 
 /** Index based on a residual quantizer. Stored vectors are
@@ -98,6 +102,58 @@ struct IndexLocalSearchQuantizer : IndexAdditiveQuantizer {
     void train(idx_t n, const float* x) override;
 };
 
+/** Index based on a product residual quantizer.
+ */
+struct IndexProductResidualQuantizer : IndexAdditiveQuantizer {
+    /// 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
+     */
+    IndexProductResidualQuantizer(
+            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_decompress);
+
+    IndexProductResidualQuantizer();
+
+    void train(idx_t n, const float* x) override;
+};
+
+/** Index based on a product local search quantizer.
+ */
+struct IndexProductLocalSearchQuantizer : IndexAdditiveQuantizer {
+    /// 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
+     */
+    IndexProductLocalSearchQuantizer(
+            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_decompress);
+
+    IndexProductLocalSearchQuantizer();
+
+    void train(idx_t n, const float* x) override;
+};
+
 /** A "virtual" index where the elements are the residual quantizer centroids.
  *
  * Intended for use as a coarse quantizer in an IndexIVF.
@@ -121,7 +177,8 @@ struct AdditiveCoarseQuantizer : 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;
     void train(idx_t n, const float* x) override;
@@ -166,7 +223,8 @@ struct ResidualCoarseQuantizer : AdditiveCoarseQuantizer {
             const float* x,
             idx_t k,
             float* distances,
-            idx_t* labels) const override;
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
 
     ResidualCoarseQuantizer();
 };