faiss 0.3.4 → 0.4.0

data/vendor/faiss/faiss/IndexIVF.h

@@ -312,11 +312,14 @@ struct IndexIVF : Index, IndexIVFInterface {
 
     /** Get a scanner for this index (store_pairs means ignore labels)
      *
-     * The default search implementation uses this to compute the distances
+     * The default search implementation uses this to compute the distances.
+     * Use sel instead of params->sel, because sel is initialized with
+     * params->sel, but may get overridden by IndexIVF's internal logic.
      */
     virtual InvertedListScanner* get_InvertedListScanner(
             bool store_pairs = false,
-            const IDSelector* sel = nullptr) const;
+            const IDSelector* sel = nullptr,
+            const IVFSearchParameters* params = nullptr) const;
 
     /** reconstruct a vector. Works only if maintain_direct_map is set to 1 or 2
      */

data/vendor/faiss/faiss/IndexIVFAdditiveQuantizer.cpp

@@ -253,7 +253,8 @@ struct AQInvertedListScannerLUT : AQInvertedListScanner {
 
 InvertedListScanner* IndexIVFAdditiveQuantizer::get_InvertedListScanner(
         bool store_pairs,
-        const IDSelector* sel) const {
+        const IDSelector* sel,
+        const IVFSearchParameters*) const {
     FAISS_THROW_IF_NOT(!sel);
     if (metric_type == METRIC_INNER_PRODUCT) {
         if (aq->search_type == AdditiveQuantizer::ST_decompress) {

data/vendor/faiss/faiss/IndexIVFAdditiveQuantizer.h

@@ -52,7 +52,8 @@ struct IndexIVFAdditiveQuantizer : IndexIVF {
 
     InvertedListScanner* get_InvertedListScanner(
             bool store_pairs,
-            const IDSelector* sel) const override;
+            const IDSelector* sel,
+            const IVFSearchParameters* params) const override;
 
     void sa_decode(idx_t n, const uint8_t* codes, float* x) const override;
 

data/vendor/faiss/faiss/IndexIVFFlat.cpp

@@ -223,7 +223,8 @@ InvertedListScanner* get_InvertedListScanner1(
 
 InvertedListScanner* IndexIVFFlat::get_InvertedListScanner(
         bool store_pairs,
-        const IDSelector* sel) const {
+        const IDSelector* sel,
+        const IVFSearchParameters*) const {
     if (sel) {
         return get_InvertedListScanner1<true>(this, store_pairs, sel);
     } else {

data/vendor/faiss/faiss/IndexIVFFlat.h

@@ -44,7 +44,8 @@ struct IndexIVFFlat : IndexIVF {
 
     InvertedListScanner* get_InvertedListScanner(
             bool store_pairs,
-            const IDSelector* sel) const override;
+            const IDSelector* sel,
+            const IVFSearchParameters* params) const override;
 
     void reconstruct_from_offset(int64_t list_no, int64_t offset, float* recons)
             const override;

data/vendor/faiss/faiss/IndexIVFPQ.cpp

@@ -1321,7 +1321,8 @@ InvertedListScanner* get_InvertedListScanner2(
 
 InvertedListScanner* IndexIVFPQ::get_InvertedListScanner(
         bool store_pairs,
-        const IDSelector* sel) const {
+        const IDSelector* sel,
+        const IVFSearchParameters*) const {
     if (sel) {
         return get_InvertedListScanner2<true>(*this, store_pairs, sel);
     } else {

data/vendor/faiss/faiss/IndexIVFPQ.h

@@ -134,7 +134,8 @@ struct IndexIVFPQ : IndexIVF {
 
     InvertedListScanner* get_InvertedListScanner(
             bool store_pairs,
-            const IDSelector* sel) const override;
+            const IDSelector* sel,
+            const IVFSearchParameters* params) const override;
 
     /// build precomputed table
     void precompute_table();

data/vendor/faiss/faiss/IndexIVFRaBitQ.cpp

@@ -0,0 +1,277 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and 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/IndexIVFRaBitQ.h>
+
+#include <omp.h>
+
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+#include <vector>
+
+#include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/RaBitQuantizer.h>
+
+namespace faiss {
+
+IndexIVFRaBitQ::IndexIVFRaBitQ(
+        Index* quantizer,
+        const size_t d,
+        const size_t nlist,
+        MetricType metric)
+        : IndexIVF(quantizer, d, nlist, 0, metric), rabitq(d, metric) {
+    code_size = rabitq.code_size;
+    invlists->code_size = code_size;
+    is_trained = false;
+
+    by_residual = true;
+}
+
+IndexIVFRaBitQ::IndexIVFRaBitQ() {
+    by_residual = true;
+}
+
+void IndexIVFRaBitQ::train_encoder(
+        idx_t n,
+        const float* x,
+        const idx_t* assign) {
+    rabitq.train(n, x);
+}
+
+void IndexIVFRaBitQ::encode_vectors(
+        idx_t n,
+        const float* x,
+        const idx_t* list_nos,
+        uint8_t* codes,
+        bool include_listnos) const {
+    size_t coarse_size = include_listnos ? coarse_code_size() : 0;
+    memset(codes, 0, (code_size + coarse_size) * n);
+
+#pragma omp parallel if (n > 1000)
+    {
+        std::vector<float> centroid(d);
+
+#pragma omp for
+        for (idx_t i = 0; i < n; i++) {
+            int64_t list_no = list_nos[i];
+            if (list_no >= 0) {
+                const float* xi = x + i * d;
+                uint8_t* code = codes + i * (code_size + coarse_size);
+
+                // both by_residual and !by_residual lead to the same code
+                quantizer->reconstruct(list_no, centroid.data());
+                rabitq.compute_codes_core(
+                        xi, code + coarse_size, 1, centroid.data());
+
+                if (coarse_size) {
+                    encode_listno(list_no, code);
+                }
+            }
+        }
+    }
+}
+
+void IndexIVFRaBitQ::add_core(
+        idx_t n,
+        const float* x,
+        const idx_t* xids,
+        const idx_t* precomputed_idx,
+        void* inverted_list_context) {
+    FAISS_THROW_IF_NOT(is_trained);
+
+    DirectMapAdd dm_add(direct_map, n, xids);
+
+#pragma omp parallel
+    {
+        std::vector<uint8_t> one_code(code_size);
+        std::vector<float> centroid(d);
+
+        int nt = omp_get_num_threads();
+        int rank = omp_get_thread_num();
+
+        // each thread takes care of a subset of lists
+        for (size_t i = 0; i < n; i++) {
+            int64_t list_no = precomputed_idx[i];
+            if (list_no >= 0 && list_no % nt == rank) {
+                int64_t id = xids ? xids[i] : ntotal + i;
+
+                const float* xi = x + i * d;
+
+                // both by_residual and !by_residual lead to the same code
+                quantizer->reconstruct(list_no, centroid.data());
+                rabitq.compute_codes_core(
+                        xi, one_code.data(), 1, centroid.data());
+
+                size_t ofs = invlists->add_entry(
+                        list_no, id, one_code.data(), inverted_list_context);
+
+                dm_add.add(i, list_no, ofs);
+
+            } else if (rank == 0 && list_no == -1) {
+                dm_add.add(i, -1, 0);
+            }
+        }
+    }
+
+    ntotal += n;
+}
+
+struct RaBitInvertedListScanner : InvertedListScanner {
+    const IndexIVFRaBitQ& ivf_rabitq;
+
+    std::vector<float> reconstructed_centroid;
+    std::vector<float> query_vector;
+
+    std::unique_ptr<FlatCodesDistanceComputer> dc;
+
+    uint8_t qb = 0;
+
+    RaBitInvertedListScanner(
+            const IndexIVFRaBitQ& ivf_rabitq_in,
+            bool store_pairs = false,
+            const IDSelector* sel = nullptr,
+            uint8_t qb_in = 0)
+            : InvertedListScanner(store_pairs, sel),
+              ivf_rabitq{ivf_rabitq_in},
+              qb{qb_in} {
+        keep_max = is_similarity_metric(ivf_rabitq.metric_type);
+        code_size = ivf_rabitq.code_size;
+    }
+
+    /// from now on we handle this query.
+    void set_query(const float* query_vector_in) override {
+        query_vector.assign(query_vector_in, query_vector_in + ivf_rabitq.d);
+
+        internal_try_setup_dc();
+    }
+
+    /// following codes come from this inverted list
+    void set_list(idx_t list_no, float coarse_dis) override {
+        this->list_no = list_no;
+
+        reconstructed_centroid.resize(ivf_rabitq.d);
+        ivf_rabitq.quantizer->reconstruct(
+                list_no, reconstructed_centroid.data());
+
+        internal_try_setup_dc();
+    }
+
+    /// compute a single query-to-code distance
+    float distance_to_code(const uint8_t* code) const override {
+        return dc->distance_to_code(code);
+    }
+
+    void internal_try_setup_dc() {
+        if (!query_vector.empty() && !reconstructed_centroid.empty()) {
+            // both query_vector and centroid are available!
+            // set up DistanceComputer
+            dc.reset(ivf_rabitq.rabitq.get_distance_computer(
+                    qb, reconstructed_centroid.data()));
+
+            dc->set_query(query_vector.data());
+        }
+    }
+};
+
+InvertedListScanner* IndexIVFRaBitQ::get_InvertedListScanner(
+        bool store_pairs,
+        const IDSelector* sel,
+        const IVFSearchParameters* search_params_in) const {
+    uint8_t used_qb = qb;
+    if (auto params = dynamic_cast<const IVFRaBitQSearchParameters*>(
+                search_params_in)) {
+        used_qb = params->qb;
+    }
+
+    return new RaBitInvertedListScanner(*this, store_pairs, sel, used_qb);
+}
+
+void IndexIVFRaBitQ::reconstruct_from_offset(
+        int64_t list_no,
+        int64_t offset,
+        float* recons) const {
+    const uint8_t* code = invlists->get_single_code(list_no, offset);
+
+    std::vector<float> centroid(d);
+    quantizer->reconstruct(list_no, centroid.data());
+
+    rabitq.decode_core(code, recons, 1, centroid.data());
+}
+
+void IndexIVFRaBitQ::sa_decode(idx_t n, const uint8_t* codes, float* x) const {
+    size_t coarse_size = coarse_code_size();
+
+#pragma omp parallel
+    {
+        std::vector<float> centroid(d);
+
+#pragma omp for
+        for (idx_t i = 0; i < n; i++) {
+            const uint8_t* code = codes + i * (code_size + coarse_size);
+            int64_t list_no = decode_listno(code);
+            float* xi = x + i * d;
+
+            quantizer->reconstruct(list_no, centroid.data());
+            rabitq.decode_core(code + coarse_size, xi, 1, centroid.data());
+        }
+    }
+}
+
+struct IVFRaBitDistanceComputer : DistanceComputer {
+    const float* q = nullptr;
+    const IndexIVFRaBitQ* parent = nullptr;
+
+    void set_query(const float* x) override;
+
+    float operator()(idx_t i) override;
+
+    float symmetric_dis(idx_t i, idx_t j) override;
+};
+
+void IVFRaBitDistanceComputer::set_query(const float* x) {
+    q = x;
+}
+
+float IVFRaBitDistanceComputer::operator()(idx_t i) {
+    // find the appropriate list
+    idx_t lo = parent->direct_map.get(i);
+    uint64_t list_no = lo_listno(lo);
+    uint64_t offset = lo_offset(lo);
+
+    const uint8_t* code = parent->invlists->get_single_code(list_no, offset);
+
+    // ok, we know the appropriate cluster that we need
+    std::vector<float> centroid(parent->d);
+    parent->quantizer->reconstruct(list_no, centroid.data());
+
+    // compute the distance
+    float distance = 0;
+
+    std::unique_ptr<FlatCodesDistanceComputer> dc(
+            parent->rabitq.get_distance_computer(parent->qb, centroid.data()));
+    dc->set_query(q);
+    distance = dc->distance_to_code(code);
+
+    // deallocate
+    parent->invlists->release_codes(list_no, code);
+
+    // done
+    return distance;
+}
+
+float IVFRaBitDistanceComputer::symmetric_dis(idx_t i, idx_t j) {
+    FAISS_THROW_MSG("Not implemented");
+}
+
+DistanceComputer* IndexIVFRaBitQ::get_distance_computer() const {
+    IVFRaBitDistanceComputer* dc = new IVFRaBitDistanceComputer;
+    dc->parent = this;
+    return dc;
+}
+
+} // namespace faiss

data/vendor/faiss/faiss/IndexIVFRaBitQ.h

@@ -0,0 +1,70 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and 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 <cstddef>
+#include <cstdint>
+
+#include <faiss/Index.h>
+#include <faiss/IndexIVF.h>
+
+#include <faiss/impl/RaBitQuantizer.h>
+
+namespace faiss {
+
+struct IVFRaBitQSearchParameters : IVFSearchParameters {
+    uint8_t qb = 0;
+};
+
+// * by_residual is true, just by design
+struct IndexIVFRaBitQ : IndexIVF {
+    RaBitQuantizer rabitq;
+
+    // the default number of bits to quantize a query with.
+    // use '0' to disable quantization and use raw fp32 values.
+    uint8_t qb = 0;
+
+    IndexIVFRaBitQ(
+            Index* quantizer,
+            const size_t d,
+            const size_t nlist,
+            MetricType metric = METRIC_L2);
+
+    IndexIVFRaBitQ();
+
+    void train_encoder(idx_t n, const float* x, const idx_t* assign) override;
+
+    void encode_vectors(
+            idx_t n,
+            const float* x,
+            const idx_t* list_nos,
+            uint8_t* codes,
+            bool include_listnos = false) const override;
+
+    void add_core(
+            idx_t n,
+            const float* x,
+            const idx_t* xids,
+            const idx_t* precomputed_idx,
+            void* inverted_list_context = nullptr) override;
+
+    InvertedListScanner* get_InvertedListScanner(
+            bool store_pairs,
+            const IDSelector* sel,
+            const IVFSearchParameters* params) const override;
+
+    void reconstruct_from_offset(int64_t list_no, int64_t offset, float* recons)
+            const override;
+
+    void sa_decode(idx_t n, const uint8_t* bytes, float* x) const override;
+
+    // unfortunately
+    DistanceComputer* get_distance_computer() const override;
+};
+
+} // namespace faiss

data/vendor/faiss/faiss/IndexIVFSpectralHash.cpp

@@ -301,7 +301,8 @@ struct BuildScanner {
 
 InvertedListScanner* IndexIVFSpectralHash::get_InvertedListScanner(
         bool store_pairs,
-        const IDSelector* sel) const {
+        const IDSelector* sel,
+        const IVFSearchParameters*) const {
     FAISS_THROW_IF_NOT(!sel);
     BuildScanner bs;
     return dispatch_HammingComputer(code_size, bs, this, store_pairs);

data/vendor/faiss/faiss/IndexIVFSpectralHash.h

@@ -71,7 +71,8 @@ struct IndexIVFSpectralHash : IndexIVF {
 
     InvertedListScanner* get_InvertedListScanner(
             bool store_pairs,
-            const IDSelector* sel) const override;
+            const IDSelector* sel,
+            const IVFSearchParameters* params) const override;
 
     /** replace the vector transform for an empty (and possibly untrained) index
      */

data/vendor/faiss/faiss/IndexRaBitQ.cpp

@@ -0,0 +1,148 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and 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/IndexRaBitQ.h>
+
+#include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/ResultHandler.h>
+
+namespace faiss {
+
+IndexRaBitQ::IndexRaBitQ() = default;
+
+IndexRaBitQ::IndexRaBitQ(idx_t d, MetricType metric)
+        : IndexFlatCodes(0, d, metric), rabitq(d, metric) {
+    code_size = rabitq.code_size;
+
+    is_trained = false;
+}
+
+void IndexRaBitQ::train(idx_t n, const float* x) {
+    // compute a centroid
+    std::vector<float> centroid(d, 0);
+    for (size_t i = 0; i < n; i++) {
+        for (size_t j = 0; j < d; j++) {
+            centroid[j] += x[i * d + j];
+        }
+    }
+
+    if (n != 0) {
+        for (size_t j = 0; j < d; j++) {
+            centroid[j] /= (float)n;
+        }
+    }
+
+    center = std::move(centroid);
+
+    //
+    rabitq.train(n, x);
+    is_trained = true;
+}
+
+void IndexRaBitQ::sa_encode(idx_t n, const float* x, uint8_t* bytes) const {
+    FAISS_THROW_IF_NOT(is_trained);
+    rabitq.compute_codes_core(x, bytes, n, center.data());
+}
+
+void IndexRaBitQ::sa_decode(idx_t n, const uint8_t* bytes, float* x) const {
+    FAISS_THROW_IF_NOT(is_trained);
+    rabitq.decode_core(bytes, x, n, center.data());
+}
+
+FlatCodesDistanceComputer* IndexRaBitQ::get_FlatCodesDistanceComputer() const {
+    FlatCodesDistanceComputer* dc =
+            rabitq.get_distance_computer(qb, center.data());
+    dc->code_size = rabitq.code_size;
+    dc->codes = codes.data();
+    return dc;
+}
+
+FlatCodesDistanceComputer* IndexRaBitQ::get_quantized_distance_computer(
+        const uint8_t qb) const {
+    FlatCodesDistanceComputer* dc =
+            rabitq.get_distance_computer(qb, center.data());
+    dc->code_size = rabitq.code_size;
+    dc->codes = codes.data();
+    return dc;
+}
+
+namespace {
+
+struct Run_search_with_dc_res {
+    using T = void;
+
+    uint8_t qb = 0;
+
+    template <class BlockResultHandler>
+    void f(BlockResultHandler& res, const IndexRaBitQ* index, const float* xq) {
+        size_t ntotal = index->ntotal;
+        using SingleResultHandler =
+                typename BlockResultHandler::SingleResultHandler;
+        const int d = index->d;
+
+#pragma omp parallel // if (res.nq > 100)
+        {
+            std::unique_ptr<FlatCodesDistanceComputer> dc(
+                    index->get_quantized_distance_computer(qb));
+            SingleResultHandler resi(res);
+#pragma omp for
+            for (int64_t q = 0; q < res.nq; q++) {
+                resi.begin(q);
+                dc->set_query(xq + d * q);
+                for (size_t i = 0; i < ntotal; i++) {
+                    if (res.is_in_selection(i)) {
+                        float dis = (*dc)(i);
+                        resi.add_result(dis, i);
+                    }
+                }
+                resi.end();
+            }
+        }
+    }
+};
+
+} // namespace
+
+void IndexRaBitQ::search(
+        idx_t n,
+        const float* x,
+        idx_t k,
+        float* distances,
+        idx_t* labels,
+        const SearchParameters* params_in) const {
+    uint8_t used_qb = qb;
+    if (auto params = dynamic_cast<const RaBitQSearchParameters*>(params_in)) {
+        used_qb = params->qb;
+    }
+
+    const IDSelector* sel = (params_in != nullptr) ? params_in->sel : nullptr;
+    Run_search_with_dc_res r;
+    r.qb = used_qb;
+
+    dispatch_knn_ResultHandler(
+            n, distances, labels, k, metric_type, sel, r, this, x);
+}
+
+void IndexRaBitQ::range_search(
+        idx_t n,
+        const float* x,
+        float radius,
+        RangeSearchResult* result,
+        const SearchParameters* params_in) const {
+    uint8_t used_qb = qb;
+    if (auto params = dynamic_cast<const RaBitQSearchParameters*>(params_in)) {
+        used_qb = params->qb;
+    }
+
+    const IDSelector* sel = (params_in != nullptr) ? params_in->sel : nullptr;
+    Run_search_with_dc_res r;
+    r.qb = used_qb;
+
+    dispatch_range_ResultHandler(result, radius, metric_type, sel, r, this, x);
+}
+
+} // namespace faiss

data/vendor/faiss/faiss/IndexRaBitQ.h

@@ -0,0 +1,65 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and 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/IndexFlatCodes.h>
+#include <faiss/impl/RaBitQuantizer.h>
+
+namespace faiss {
+
+struct RaBitQSearchParameters : SearchParameters {
+    uint8_t qb = 0;
+};
+
+struct IndexRaBitQ : IndexFlatCodes {
+    RaBitQuantizer rabitq;
+
+    // center of all points
+    std::vector<float> center;
+
+    // the default number of bits to quantize a query with.
+    // use '0' to disable quantization and use raw fp32 values.
+    uint8_t qb = 0;
+
+    IndexRaBitQ();
+
+    IndexRaBitQ(idx_t d, MetricType metric = METRIC_L2);
+
+    void train(idx_t n, const float* x) 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;
+
+    // returns a quantized-to-qb bits DC if qb > 0
+    // returns a default fp32-based DC if qb == 0
+    FlatCodesDistanceComputer* get_FlatCodesDistanceComputer() const override;
+
+    // returns a quantized-to-qb bits DC if qb_in > 0
+    // returns a default fp32-based DC if qb_in == 0
+    FlatCodesDistanceComputer* get_quantized_distance_computer(
+            const uint8_t qb_in) const;
+
+    // Don't rely on sa_decode(), bcz it is good for IP, but not for L2.
+    //   As a result, use get_FlatCodesDistanceComputer() for the search.
+    void search(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
+
+    void range_search(
+            idx_t n,
+            const float* x,
+            float radius,
+            RangeSearchResult* result,
+            const SearchParameters* params = nullptr) const override;
+};
+
+} // namespace faiss

data/vendor/faiss/faiss/IndexScalarQuantizer.cpp

@@ -254,7 +254,8 @@ void IndexIVFScalarQuantizer::add_core(
 
 InvertedListScanner* IndexIVFScalarQuantizer::get_InvertedListScanner(
         bool store_pairs,
-        const IDSelector* sel) const {
+        const IDSelector* sel,
+        const IVFSearchParameters*) const {
     return sq.select_InvertedListScanner(
             metric_type, quantizer, store_pairs, sel, by_residual);
 }

data/vendor/faiss/faiss/IndexScalarQuantizer.h

@@ -96,7 +96,8 @@ struct IndexIVFScalarQuantizer : IndexIVF {
 
     InvertedListScanner* get_InvertedListScanner(
             bool store_pairs,
-            const IDSelector* sel) const override;
+            const IDSelector* sel,
+            const IVFSearchParameters* params) const override;
 
     void reconstruct_from_offset(int64_t list_no, int64_t offset, float* recons)
             const override;