faiss 0.3.4 → 0.4.1

data/vendor/faiss/faiss/IndexHNSW.cpp

@@ -8,9 +8,7 @@
 #include <faiss/IndexHNSW.h>
 
 #include <omp.h>
-#include <cassert>
 #include <cinttypes>
-#include <cmath>
 #include <cstdio>
 #include <cstdlib>
 #include <cstring>
@@ -124,7 +122,7 @@ void hnsw_add_vertices(
         int i1 = n;
 
         for (int pt_level = hist.size() - 1;
-             pt_level >= !index_hnsw.init_level0;
+             pt_level >= int(!index_hnsw.init_level0);
              pt_level--) {
             int i0 = i1 - hist[pt_level];
 
@@ -212,7 +210,9 @@ IndexHNSW::IndexHNSW(int d, int M, MetricType metric)
         : Index(d, metric), hnsw(M) {}
 
 IndexHNSW::IndexHNSW(Index* storage, int M)
-        : Index(storage->d, storage->metric_type), hnsw(M), storage(storage) {}
+        : Index(storage->d, storage->metric_type), hnsw(M), storage(storage) {
+    metric_arg = storage->metric_arg;
+}
 
 IndexHNSW::~IndexHNSW() {
     if (own_fields) {
@@ -237,19 +237,19 @@ void hnsw_search(
         idx_t n,
         const float* x,
         BlockResultHandler& bres,
-        const SearchParameters* params_in) {
+        const SearchParameters* params) {
     FAISS_THROW_IF_NOT_MSG(
             index->storage,
             "No storage index, please use IndexHNSWFlat (or variants) "
             "instead of IndexHNSW directly");
-    const SearchParametersHNSW* params = nullptr;
     const HNSW& hnsw = index->hnsw;
 
     int efSearch = hnsw.efSearch;
-    if (params_in) {
-        params = dynamic_cast<const SearchParametersHNSW*>(params_in);
-        FAISS_THROW_IF_NOT_MSG(params, "params type invalid");
-        efSearch = params->efSearch;
+    if (params) {
+        if (const SearchParametersHNSW* hnsw_params =
+                    dynamic_cast<const SearchParametersHNSW*>(params)) {
+            efSearch = hnsw_params->efSearch;
+        }
     }
     size_t n1 = 0, n2 = 0, ndis = 0, nhops = 0;
 
@@ -294,13 +294,13 @@ void IndexHNSW::search(
         idx_t k,
         float* distances,
         idx_t* labels,
-        const SearchParameters* params_in) const {
+        const SearchParameters* params) const {
     FAISS_THROW_IF_NOT(k > 0);
 
     using RH = HeapBlockResultHandler<HNSW::C>;
     RH bres(n, distances, labels, k);
 
-    hnsw_search(this, n, x, bres, params_in);
+    hnsw_search(this, n, x, bres, params);
 
     if (is_similarity_metric(this->metric_type)) {
         // we need to revert the negated distances
@@ -408,17 +408,10 @@ void IndexHNSW::search_level_0(
         idx_t* labels,
         int nprobe,
         int search_type,
-        const SearchParameters* params_in) const {
+        const SearchParameters* params) const {
     FAISS_THROW_IF_NOT(k > 0);
     FAISS_THROW_IF_NOT(nprobe > 0);
 
-    const SearchParametersHNSW* params = nullptr;
-
-    if (params_in) {
-        params = dynamic_cast<const SearchParametersHNSW*>(params_in);
-        FAISS_THROW_IF_NOT_MSG(params, "params type invalid");
-    }
-
     storage_idx_t ntotal = hnsw.levels.size();
 
     using RH = HeapBlockResultHandler<HNSW::C>;

data/vendor/faiss/faiss/IndexHNSW.h

@@ -138,7 +138,7 @@ struct IndexHNSWPQ : IndexHNSW {
     void train(idx_t n, const float* x) override;
 };
 
-/** SQ index topped with with a HNSW structure to access elements
+/** SQ index topped with a HNSW structure to access elements
  *  more efficiently.
  */
 struct IndexHNSWSQ : IndexHNSW {

data/vendor/faiss/faiss/IndexIVF.cpp

@@ -455,7 +455,7 @@ void IndexIVF::search_preassigned(
 #pragma omp parallel if (do_parallel) reduction(+ : nlistv, ndis, nheap)
     {
         std::unique_ptr<InvertedListScanner> scanner(
-                get_InvertedListScanner(store_pairs, sel));
+                get_InvertedListScanner(store_pairs, sel, params));
 
         /*****************************************************
          * Depending on parallel_mode, there are two possible ways
@@ -796,7 +796,7 @@ void IndexIVF::range_search_preassigned(
     {
         RangeSearchPartialResult pres(result);
         std::unique_ptr<InvertedListScanner> scanner(
-                get_InvertedListScanner(store_pairs, sel));
+                get_InvertedListScanner(store_pairs, sel, params));
         FAISS_THROW_IF_NOT(scanner.get());
         all_pres[omp_get_thread_num()] = &pres;
 
@@ -912,7 +912,8 @@ void IndexIVF::range_search_preassigned(
 
 InvertedListScanner* IndexIVF::get_InvertedListScanner(
         bool /*store_pairs*/,
-        const IDSelector* /* sel */) const {
+        const IDSelector* /* sel */,
+        const IVFSearchParameters* /* params */) const {
     FAISS_THROW_MSG("get_InvertedListScanner not implemented");
 }
 
@@ -1290,6 +1291,14 @@ size_t InvertedListScanner::scan_codes(
 
     if (!keep_max) {
         for (size_t j = 0; j < list_size; j++) {
+            if (sel != nullptr) {
+                int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
+                if (!sel->is_member(id)) {
+                    codes += code_size;
+                    continue;
+                }
+            }
+
             float dis = distance_to_code(codes);
             if (dis < simi[0]) {
                 int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
@@ -1300,6 +1309,14 @@ size_t InvertedListScanner::scan_codes(
         }
     } else {
         for (size_t j = 0; j < list_size; j++) {
+            if (sel != nullptr) {
+                int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
+                if (!sel->is_member(id)) {
+                    codes += code_size;
+                    continue;
+                }
+            }
+
             float dis = distance_to_code(codes);
             if (dis > simi[0]) {
                 int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];

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
      */