faiss 0.2.3 → 0.2.5

data/vendor/faiss/faiss/IndexAdditiveQuantizer.h

@@ -0,0 +1,253 @@
+/**
+ * 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.
+ */
+
+#ifndef FAISS_INDEX_ADDITIVE_QUANTIZER_H
+#define FAISS_INDEX_ADDITIVE_QUANTIZER_H
+
+#include <faiss/impl/AdditiveQuantizer.h>
+
+#include <cstdint>
+#include <vector>
+
+#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>
+
+namespace faiss {
+
+/// Abstract class for additive quantizers. The search functions are in common.
+struct IndexAdditiveQuantizer : IndexFlatCodes {
+    // the quantizer, this points to the relevant field in the inheriting
+    // classes
+    AdditiveQuantizer* aq;
+    using Search_type_t = AdditiveQuantizer::Search_type_t;
+
+    explicit IndexAdditiveQuantizer(
+            idx_t d,
+            AdditiveQuantizer* aq,
+            MetricType metric = METRIC_L2);
+
+    void search(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            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
+ * approximated by residual quantization codes.
+ * Can also be used as a codec
+ */
+struct IndexResidualQuantizer : IndexAdditiveQuantizer {
+    /// 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
+     */
+    IndexResidualQuantizer(
+            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_decompress);
+
+    IndexResidualQuantizer(
+            int d,
+            const std::vector<size_t>& nbits,
+            MetricType metric = METRIC_L2,
+            Search_type_t search_type = AdditiveQuantizer::ST_decompress);
+
+    IndexResidualQuantizer();
+
+    void train(idx_t n, const float* x) override;
+};
+
+struct IndexLocalSearchQuantizer : IndexAdditiveQuantizer {
+    LocalSearchQuantizer lsq;
+
+    /** Constructor.
+     *
+     * @param d      dimensionality of the input vectors
+     * @param M      number of subquantizers
+     * @param nbits  number of bit per subvector index
+     */
+    IndexLocalSearchQuantizer(
+            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_decompress);
+
+    IndexLocalSearchQuantizer();
+
+    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.
+ */
+struct AdditiveCoarseQuantizer : Index {
+    AdditiveQuantizer* aq;
+
+    explicit AdditiveCoarseQuantizer(
+            idx_t d = 0,
+            AdditiveQuantizer* aq = nullptr,
+            MetricType metric = METRIC_L2);
+
+    /// norms of centroids, useful for knn-search
+    std::vector<float> centroid_norms;
+
+    /// N/A
+    void add(idx_t n, const float* x) override;
+
+    void search(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            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;
+
+    /// N/A
+    void reset() override;
+};
+
+/** The ResidualCoarseQuantizer is a bit specialized compared to the
+ * default AdditiveCoarseQuantizer because it can use a beam search
+ * at search time (slow but may be useful for very large vocabularies) */
+struct ResidualCoarseQuantizer : AdditiveCoarseQuantizer {
+    /// The residual quantizer used to encode the vectors
+    ResidualQuantizer rq;
+
+    /// factor between the beam size and the search k
+    /// if negative, use exact search-to-centroid
+    float beam_factor;
+
+    /// computes centroid norms if required
+    void set_beam_factor(float new_beam_factor);
+
+    /** Constructor.
+     *
+     * @param d      dimensionality of the input vectors
+     * @param M      number of subquantizers
+     * @param nbits  number of bit per subvector index
+     */
+    ResidualCoarseQuantizer(
+            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);
+
+    ResidualCoarseQuantizer(
+            int d,
+            const std::vector<size_t>& nbits,
+            MetricType metric = METRIC_L2);
+
+    void search(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
+
+    ResidualCoarseQuantizer();
+};
+
+struct LocalSearchCoarseQuantizer : AdditiveCoarseQuantizer {
+    /// The residual quantizer used to encode the vectors
+    LocalSearchQuantizer lsq;
+
+    /** Constructor.
+     *
+     * @param d      dimensionality of the input vectors
+     * @param M      number of subquantizers
+     * @param nbits  number of bit per subvector index
+     */
+    LocalSearchCoarseQuantizer(
+            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);
+
+    LocalSearchCoarseQuantizer();
+};
+
+} // namespace faiss
+
+#endif

data/vendor/faiss/faiss/IndexAdditiveQuantizerFastScan.cpp

@@ -0,0 +1,299 @@
+/**
+ * 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/IndexAdditiveQuantizerFastScan.h>
+
+#include <limits.h>
+#include <cassert>
+#include <memory>
+
+#include <omp.h>
+
+#include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/LocalSearchQuantizer.h>
+#include <faiss/impl/LookupTableScaler.h>
+#include <faiss/impl/ResidualQuantizer.h>
+#include <faiss/impl/pq4_fast_scan.h>
+#include <faiss/utils/quantize_lut.h>
+#include <faiss/utils/utils.h>
+
+namespace faiss {
+
+inline size_t roundup(size_t a, size_t b) {
+    return (a + b - 1) / b * b;
+}
+
+IndexAdditiveQuantizerFastScan::IndexAdditiveQuantizerFastScan(
+        AdditiveQuantizer* aq,
+        MetricType metric,
+        int bbs) {
+    init(aq, metric, bbs);
+}
+
+void IndexAdditiveQuantizerFastScan::init(
+        AdditiveQuantizer* aq,
+        MetricType metric,
+        int bbs) {
+    FAISS_THROW_IF_NOT(aq != nullptr);
+    FAISS_THROW_IF_NOT(!aq->nbits.empty());
+    FAISS_THROW_IF_NOT(aq->nbits[0] == 4);
+    if (metric == METRIC_INNER_PRODUCT) {
+        FAISS_THROW_IF_NOT_MSG(
+                aq->search_type == AdditiveQuantizer::ST_LUT_nonorm,
+                "Search type must be ST_LUT_nonorm for IP metric");
+    } else {
+        FAISS_THROW_IF_NOT_MSG(
+                aq->search_type == AdditiveQuantizer::ST_norm_lsq2x4 ||
+                        aq->search_type == AdditiveQuantizer::ST_norm_rq2x4,
+                "Search type must be lsq2x4 or rq2x4 for L2 metric");
+    }
+
+    this->aq = aq;
+    if (metric == METRIC_L2) {
+        M = aq->M + 2; // 2x4 bits AQ
+    } else {
+        M = aq->M;
+    }
+    init_fastscan(aq->d, M, 4, metric, bbs);
+
+    max_train_points = 1024 * ksub * M;
+}
+
+IndexAdditiveQuantizerFastScan::IndexAdditiveQuantizerFastScan()
+        : IndexFastScan() {
+    is_trained = false;
+    aq = nullptr;
+}
+
+IndexAdditiveQuantizerFastScan::IndexAdditiveQuantizerFastScan(
+        const IndexAdditiveQuantizer& orig,
+        int bbs) {
+    init(orig.aq, orig.metric_type, bbs);
+
+    ntotal = orig.ntotal;
+    is_trained = orig.is_trained;
+    orig_codes = orig.codes.data();
+
+    ntotal2 = roundup(ntotal, bbs);
+    codes.resize(ntotal2 * M2 / 2);
+    pq4_pack_codes(orig_codes, ntotal, M, ntotal2, bbs, M2, codes.get());
+}
+
+IndexAdditiveQuantizerFastScan::~IndexAdditiveQuantizerFastScan() {}
+
+void IndexAdditiveQuantizerFastScan::train(idx_t n, const float* x_in) {
+    if (is_trained) {
+        return;
+    }
+
+    const int seed = 0x12345;
+    size_t nt = n;
+    const float* x = fvecs_maybe_subsample(
+            d, &nt, max_train_points, x_in, verbose, seed);
+    n = nt;
+    if (verbose) {
+        printf("training additive quantizer on %zd vectors\n", nt);
+    }
+
+    aq->verbose = verbose;
+    aq->train(n, x);
+    if (metric_type == METRIC_L2) {
+        estimate_norm_scale(n, x);
+    }
+
+    is_trained = true;
+}
+
+void IndexAdditiveQuantizerFastScan::estimate_norm_scale(
+        idx_t n,
+        const float* x_in) {
+    FAISS_THROW_IF_NOT(metric_type == METRIC_L2);
+
+    constexpr int seed = 0x980903;
+    constexpr size_t max_points_estimated = 65536;
+    size_t ns = n;
+    const float* x = fvecs_maybe_subsample(
+            d, &ns, max_points_estimated, x_in, verbose, seed);
+    n = ns;
+    std::unique_ptr<float[]> del_x;
+    if (x != x_in) {
+        del_x.reset((float*)x);
+    }
+
+    std::vector<float> dis_tables(n * M * ksub);
+    compute_float_LUT(dis_tables.data(), n, x);
+
+    // here we compute the mean of scales for each query
+    // TODO: try max of scales
+    double scale = 0;
+
+#pragma omp parallel for reduction(+ : scale)
+    for (idx_t i = 0; i < n; i++) {
+        const float* lut = dis_tables.data() + i * M * ksub;
+        scale += quantize_lut::aq_estimate_norm_scale(M, ksub, 2, lut);
+    }
+    scale /= n;
+    norm_scale = (int)std::roundf(std::max(scale, 1.0));
+
+    if (verbose) {
+        printf("estimated norm scale: %lf\n", scale);
+        printf("rounded norm scale: %d\n", norm_scale);
+    }
+}
+
+void IndexAdditiveQuantizerFastScan::compute_codes(
+        uint8_t* tmp_codes,
+        idx_t n,
+        const float* x) const {
+    aq->compute_codes(x, tmp_codes, n);
+}
+
+void IndexAdditiveQuantizerFastScan::compute_float_LUT(
+        float* lut,
+        idx_t n,
+        const float* x) const {
+    if (metric_type == METRIC_INNER_PRODUCT) {
+        aq->compute_LUT(n, x, lut, 1.0f);
+    } else {
+        // compute inner product look-up tables
+        const size_t ip_dim12 = aq->M * ksub;
+        const size_t norm_dim12 = 2 * ksub;
+        std::vector<float> ip_lut(n * ip_dim12);
+        aq->compute_LUT(n, x, ip_lut.data(), -2.0f);
+
+        // copy and rescale norm look-up tables
+        auto norm_tabs = aq->norm_tabs;
+        if (rescale_norm && norm_scale > 1 && metric_type == METRIC_L2) {
+            for (size_t i = 0; i < norm_tabs.size(); i++) {
+                norm_tabs[i] /= norm_scale;
+            }
+        }
+        const float* norm_lut = norm_tabs.data();
+        FAISS_THROW_IF_NOT(norm_tabs.size() == norm_dim12);
+
+        // combine them
+        for (idx_t i = 0; i < n; i++) {
+            memcpy(lut, ip_lut.data() + i * ip_dim12, ip_dim12 * sizeof(*lut));
+            lut += ip_dim12;
+            memcpy(lut, norm_lut, norm_dim12 * sizeof(*lut));
+            lut += norm_dim12;
+        }
+    }
+}
+
+void IndexAdditiveQuantizerFastScan::search(
+        idx_t n,
+        const float* x,
+        idx_t k,
+        float* distances,
+        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);
+    bool rescale = (rescale_norm && norm_scale > 1 && metric_type == METRIC_L2);
+    if (!rescale) {
+        IndexFastScan::search(n, x, k, distances, labels);
+        return;
+    }
+
+    NormTableScaler scaler(norm_scale);
+    if (metric_type == METRIC_L2) {
+        search_dispatch_implem<true>(n, x, k, distances, labels, scaler);
+    } else {
+        search_dispatch_implem<false>(n, x, k, distances, labels, scaler);
+    }
+}
+
+void IndexAdditiveQuantizerFastScan::sa_decode(
+        idx_t n,
+        const uint8_t* bytes,
+        float* x) const {
+    aq->decode(bytes, x, n);
+}
+
+/**************************************************************************************
+ * IndexResidualQuantizerFastScan
+ **************************************************************************************/
+
+IndexResidualQuantizerFastScan::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,
+        Search_type_t search_type,
+        int bbs)
+        : rq(d, M, nbits, search_type) {
+    init(&rq, metric, bbs);
+}
+
+IndexResidualQuantizerFastScan::IndexResidualQuantizerFastScan() {
+    aq = &rq;
+}
+
+/**************************************************************************************
+ * IndexLocalSearchQuantizerFastScan
+ **************************************************************************************/
+
+IndexLocalSearchQuantizerFastScan::IndexLocalSearchQuantizerFastScan(
+        int d,
+        size_t M,     ///< number of subquantizers
+        size_t nbits, ///< number of bit per subvector index
+        MetricType metric,
+        Search_type_t search_type,
+        int bbs)
+        : lsq(d, M, nbits, search_type) {
+    init(&lsq, metric, bbs);
+}
+
+IndexLocalSearchQuantizerFastScan::IndexLocalSearchQuantizerFastScan() {
+    aq = &lsq;
+}
+
+/**************************************************************************************
+ * IndexProductResidualQuantizerFastScan
+ **************************************************************************************/
+
+IndexProductResidualQuantizerFastScan::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,
+        Search_type_t search_type,
+        int bbs)
+        : prq(d, nsplits, Msub, nbits, search_type) {
+    init(&prq, metric, bbs);
+}
+
+IndexProductResidualQuantizerFastScan::IndexProductResidualQuantizerFastScan() {
+    aq = &prq;
+}
+
+/**************************************************************************************
+ * IndexProductLocalSearchQuantizerFastScan
+ **************************************************************************************/
+
+IndexProductLocalSearchQuantizerFastScan::
+        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,
+                Search_type_t search_type,
+                int bbs)
+        : plsq(d, nsplits, Msub, nbits, search_type) {
+    init(&plsq, metric, bbs);
+}
+
+IndexProductLocalSearchQuantizerFastScan::
+        IndexProductLocalSearchQuantizerFastScan() {
+    aq = &plsq;
+}
+
+} // namespace faiss