faiss 0.3.0 → 0.3.1

data/vendor/faiss/faiss/IndexIVFFastScan.h

@@ -14,6 +14,9 @@
 
 namespace faiss {
 
+struct NormTableScaler;
+struct SIMDResultHandlerToFloat;
+
 /** Fast scan version of IVFPQ and IVFAQ. Works for 4-bit PQ/AQ for now.
  *
  * The codes in the inverted lists are not stored sequentially but
@@ -28,6 +31,12 @@ namespace faiss {
  * 11: idem, collect results in reservoir
  * 12: optimizer int16 search, collect results in heap, uses qbs
  * 13: idem, collect results in reservoir
+ * 14: internally multithreaded implem over nq * nprobe
+ * 15: same with reservoir
+ *
+ * For range search, only 10 and 12 are supported.
+ * add 100 to the implem to force single-thread scanning (the coarse quantizer
+ * may still use multiple threads).
  */
 
 struct IndexIVFFastScan : IndexIVF {
@@ -45,7 +54,6 @@ struct IndexIVFFastScan : IndexIVF {
     int implem = 0;
     // skip some parts of the computation (for timing)
     int skip = 0;
-    bool by_residual = false;
 
     // batching factors at search time (0 = default)
     int qbs = 0;
@@ -81,19 +89,24 @@ struct IndexIVFFastScan : IndexIVF {
 
     virtual bool lookup_table_is_3d() const = 0;
 
+    // compact way of conveying coarse quantization results
+    struct CoarseQuantized {
+        size_t nprobe;
+        const float* dis = nullptr;
+        const idx_t* ids = nullptr;
+    };
+
     virtual void compute_LUT(
             size_t n,
             const float* x,
-            const idx_t* coarse_ids,
-            const float* coarse_dis,
+            const CoarseQuantized& cq,
             AlignedTable<float>& dis_tables,
             AlignedTable<float>& biases) const = 0;
 
     void compute_LUT_uint8(
             size_t n,
             const float* x,
-            const idx_t* coarse_ids,
-            const float* coarse_dis,
+            const CoarseQuantized& cq,
             AlignedTable<uint8_t>& dis_tables,
             AlignedTable<uint16_t>& biases,
             float* normalizers) const;
@@ -106,7 +119,18 @@ struct IndexIVFFastScan : IndexIVF {
             idx_t* labels,
             const SearchParameters* params = nullptr) const override;
 
-    /// will just fail
+    void search_preassigned(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            const idx_t* assign,
+            const float* centroid_dis,
+            float* distances,
+            idx_t* labels,
+            bool store_pairs,
+            const IVFSearchParameters* params = nullptr,
+            IndexIVFStats* stats = nullptr) const override;
+
     void range_search(
             idx_t n,
             const float* x,
@@ -116,69 +140,75 @@ struct IndexIVFFastScan : IndexIVF {
 
     // internal search funcs
 
-    template <bool is_max, class Scaler>
+    // dispatch to implementations and parallelize
     void search_dispatch_implem(
             idx_t n,
             const float* x,
             idx_t k,
             float* distances,
             idx_t* labels,
-            const Scaler& scaler) const;
+            const CoarseQuantized& cq,
+            const NormTableScaler* scaler) const;
 
-    template <class C, class Scaler>
+    void range_search_dispatch_implem(
+            idx_t n,
+            const float* x,
+            float radius,
+            RangeSearchResult& rres,
+            const CoarseQuantized& cq_in,
+            const NormTableScaler* scaler) const;
+
+    // impl 1 and 2 are just for verification
+    template <class C>
     void search_implem_1(
             idx_t n,
             const float* x,
             idx_t k,
             float* distances,
             idx_t* labels,
-            const Scaler& scaler) const;
+            const CoarseQuantized& cq,
+            const NormTableScaler* scaler) const;
 
-    template <class C, class Scaler>
+    template <class C>
     void search_implem_2(
             idx_t n,
             const float* x,
             idx_t k,
             float* distances,
             idx_t* labels,
-            const Scaler& scaler) const;
+            const CoarseQuantized& cq,
+            const NormTableScaler* scaler) const;
 
     // implem 10 and 12 are not multithreaded internally, so
     // export search stats
-    template <class C, class Scaler>
     void search_implem_10(
             idx_t n,
             const float* x,
-            idx_t k,
-            float* distances,
-            idx_t* labels,
-            int impl,
+            SIMDResultHandlerToFloat& handler,
+            const CoarseQuantized& cq,
             size_t* ndis_out,
             size_t* nlist_out,
-            const Scaler& scaler) const;
+            const NormTableScaler* scaler) const;
 
-    template <class C, class Scaler>
     void search_implem_12(
             idx_t n,
             const float* x,
-            idx_t k,
-            float* distances,
-            idx_t* labels,
-            int impl,
+            SIMDResultHandlerToFloat& handler,
+            const CoarseQuantized& cq,
             size_t* ndis_out,
             size_t* nlist_out,
-            const Scaler& scaler) const;
+            const NormTableScaler* scaler) const;
 
     // implem 14 is multithreaded internally across nprobes and queries
-    template <class C, class Scaler>
     void search_implem_14(
             idx_t n,
             const float* x,
             idx_t k,
             float* distances,
             idx_t* labels,
+            const CoarseQuantized& cq,
             int impl,
-            const Scaler& scaler) const;
+            const NormTableScaler* scaler) const;
 
     // reconstruct vectors from packed invlists
     void reconstruct_from_offset(int64_t list_no, int64_t offset, float* recons)

data/vendor/faiss/faiss/IndexIVFFlat.cpp

@@ -36,15 +36,22 @@ IndexIVFFlat::IndexIVFFlat(
         MetricType metric)
         : IndexIVF(quantizer, d, nlist, sizeof(float) * d, metric) {
     code_size = sizeof(float) * d;
+    by_residual = false;
+}
+
+IndexIVFFlat::IndexIVFFlat() {
+    by_residual = false;
 }
 
 void IndexIVFFlat::add_core(
         idx_t n,
         const float* x,
-        const int64_t* xids,
-        const int64_t* coarse_idx) {
+        const idx_t* xids,
+        const idx_t* coarse_idx,
+        void* inverted_list_context) {
     FAISS_THROW_IF_NOT(is_trained);
     FAISS_THROW_IF_NOT(coarse_idx);
+    FAISS_THROW_IF_NOT(!by_residual);
     assert(invlists);
     direct_map.check_can_add(xids);
 
@@ -64,8 +71,8 @@ void IndexIVFFlat::add_core(
             if (list_no >= 0 && list_no % nt == rank) {
                 idx_t id = xids ? xids[i] : ntotal + i;
                 const float* xi = x + i * d;
-                size_t offset =
-                        invlists->add_entry(list_no, id, (const uint8_t*)xi);
+                size_t offset = invlists->add_entry(
+                        list_no, id, (const uint8_t*)xi, inverted_list_context);
                 dm_adder.add(i, list_no, offset);
                 n_add++;
             } else if (rank == 0 && list_no == -1) {
@@ -89,6 +96,7 @@ void IndexIVFFlat::encode_vectors(
         const idx_t* list_nos,
         uint8_t* codes,
         bool include_listnos) const {
+    FAISS_THROW_IF_NOT(!by_residual);
     if (!include_listnos) {
         memcpy(codes, x, code_size * n);
     } else {
@@ -123,7 +131,9 @@ struct IVFFlatScanner : InvertedListScanner {
     size_t d;
 
     IVFFlatScanner(size_t d, bool store_pairs, const IDSelector* sel)
-            : InvertedListScanner(store_pairs, sel), d(d) {}
+            : InvertedListScanner(store_pairs, sel), d(d) {
+        keep_max = is_similarity_metric(metric);
+    }
 
     const float* xi;
     void set_query(const float* query) override {

data/vendor/faiss/faiss/IndexIVFFlat.h

@@ -32,7 +32,8 @@ struct IndexIVFFlat : IndexIVF {
             idx_t n,
             const float* x,
             const idx_t* xids,
-            const idx_t* precomputed_idx) override;
+            const idx_t* precomputed_idx,
+            void* inverted_list_context = nullptr) override;
 
     void encode_vectors(
             idx_t n,
@@ -50,7 +51,7 @@ struct IndexIVFFlat : IndexIVF {
 
     void sa_decode(idx_t n, const uint8_t* bytes, float* x) const override;
 
-    IndexIVFFlat() {}
+    IndexIVFFlat();
 };
 
 struct IndexIVFFlatDedup : IndexIVFFlat {

data/vendor/faiss/faiss/IndexIVFIndependentQuantizer.cpp

@@ -0,0 +1,172 @@
+/**
+ * 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/IndexIVFIndependentQuantizer.h>
+#include <faiss/IndexIVFPQ.h>
+#include <faiss/impl/FaissAssert.h>
+#include <faiss/utils/utils.h>
+
+namespace faiss {
+
+IndexIVFIndependentQuantizer::IndexIVFIndependentQuantizer(
+        Index* quantizer,
+        IndexIVF* index_ivf,
+        VectorTransform* vt)
+        : Index(quantizer->d, index_ivf->metric_type),
+          quantizer(quantizer),
+          vt(vt),
+          index_ivf(index_ivf) {
+    if (vt) {
+        FAISS_THROW_IF_NOT_MSG(
+                vt->d_in == d && vt->d_out == index_ivf->d,
+                "invalid vector dimensions");
+    } else {
+        FAISS_THROW_IF_NOT_MSG(index_ivf->d == d, "invalid vector dimensions");
+    }
+
+    if (quantizer->is_trained && quantizer->ntotal != 0) {
+        FAISS_THROW_IF_NOT(quantizer->ntotal == index_ivf->nlist);
+    }
+    if (index_ivf->is_trained && vt) {
+        FAISS_THROW_IF_NOT(vt->is_trained);
+    }
+    ntotal = index_ivf->ntotal;
+    is_trained =
+            (quantizer->is_trained && quantizer->ntotal == index_ivf->nlist &&
+             (!vt || vt->is_trained) && index_ivf->is_trained);
+
+    // disable precomputed tables because they use the distances that are
+    // provided by the coarse quantizer (that are out of sync with the IVFPQ)
+    if (auto index_ivfpq = dynamic_cast<IndexIVFPQ*>(index_ivf)) {
+        index_ivfpq->use_precomputed_table = -1;
+    }
+}
+
+IndexIVFIndependentQuantizer::~IndexIVFIndependentQuantizer() {
+    if (own_fields) {
+        delete quantizer;
+        delete index_ivf;
+        delete vt;
+    }
+}
+
+namespace {
+
+struct VTransformedVectors : TransformedVectors {
+    VTransformedVectors(const VectorTransform* vt, idx_t n, const float* x)
+            : TransformedVectors(x, vt ? vt->apply(n, x) : x) {}
+};
+
+struct SubsampledVectors : TransformedVectors {
+    SubsampledVectors(int d, idx_t* n, idx_t max_n, const float* x)
+            : TransformedVectors(
+                      x,
+                      fvecs_maybe_subsample(d, (size_t*)n, max_n, x, true)) {}
+};
+
+} // anonymous namespace
+
+void IndexIVFIndependentQuantizer::add(idx_t n, const float* x) {
+    std::vector<float> D(n);
+    std::vector<idx_t> I(n);
+    quantizer->search(n, x, 1, D.data(), I.data());
+
+    VTransformedVectors tv(vt, n, x);
+
+    index_ivf->add_core(n, tv.x, nullptr, I.data());
+}
+
+void IndexIVFIndependentQuantizer::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 parameters not supported");
+    int nprobe = index_ivf->nprobe;
+    std::vector<float> D(n * nprobe);
+    std::vector<idx_t> I(n * nprobe);
+    quantizer->search(n, x, nprobe, D.data(), I.data());
+
+    VTransformedVectors tv(vt, n, x);
+
+    index_ivf->search_preassigned(
+            n, tv.x, k, I.data(), D.data(), distances, labels, false);
+}
+
+void IndexIVFIndependentQuantizer::reset() {
+    index_ivf->reset();
+    ntotal = 0;
+}
+
+void IndexIVFIndependentQuantizer::train(idx_t n, const float* x) {
+    // quantizer training
+    size_t nlist = index_ivf->nlist;
+    Level1Quantizer l1(quantizer, nlist);
+    l1.train_q1(n, x, verbose, metric_type);
+
+    // train the VectorTransform
+    if (vt && !vt->is_trained) {
+        if (verbose) {
+            printf("IndexIVFIndependentQuantizer: train the VectorTransform\n");
+        }
+        vt->train(n, x);
+    }
+
+    // get the centroids from the quantizer, transform them and
+    // add them to the index_ivf's quantizer
+    if (verbose) {
+        printf("IndexIVFIndependentQuantizer: extract the main quantizer centroids\n");
+    }
+    std::vector<float> centroids(nlist * d);
+    quantizer->reconstruct_n(0, nlist, centroids.data());
+    VTransformedVectors tcent(vt, nlist, centroids.data());
+
+    if (verbose) {
+        printf("IndexIVFIndependentQuantizer: add centroids to the secondary quantizer\n");
+    }
+    if (!index_ivf->quantizer->is_trained) {
+        index_ivf->quantizer->train(nlist, tcent.x);
+    }
+    index_ivf->quantizer->add(nlist, tcent.x);
+
+    // train the payload
+
+    // optional subsampling
+    idx_t max_nt = index_ivf->train_encoder_num_vectors();
+    if (max_nt <= 0) {
+        max_nt = (size_t)1 << 35;
+    }
+    SubsampledVectors sv(index_ivf->d, &n, max_nt, x);
+
+    // transform subsampled vectors
+    VTransformedVectors tv(vt, n, sv.x);
+
+    if (verbose) {
+        printf("IndexIVFIndependentQuantizer: train encoder\n");
+    }
+
+    if (index_ivf->by_residual) {
+        // assign with quantizer
+        std::vector<idx_t> assign(n);
+        quantizer->assign(n, sv.x, assign.data());
+
+        // compute residual with IVF quantizer
+        std::vector<float> residuals(n * index_ivf->d);
+        index_ivf->quantizer->compute_residual_n(
+                n, tv.x, residuals.data(), assign.data());
+
+        index_ivf->train_encoder(n, residuals.data(), assign.data());
+    } else {
+        index_ivf->train_encoder(n, tv.x, nullptr);
+    }
+    index_ivf->is_trained = true;
+    is_trained = true;
+}
+
+} // namespace faiss

data/vendor/faiss/faiss/IndexIVFIndependentQuantizer.h

@@ -0,0 +1,56 @@
+/**
+ * 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/IndexIVF.h>
+#include <faiss/VectorTransform.h>
+
+namespace faiss {
+
+/** An IVF index with a quantizer that has a different input dimension from the
+ * payload size. The vectors to encode are obtained from the input vectors by a
+ * VectorTransform.
+ */
+struct IndexIVFIndependentQuantizer : Index {
+    /// quantizer is fed directly with the input vectors
+    Index* quantizer = nullptr;
+
+    /// transform before the IVF vectors are applied
+    VectorTransform* vt = nullptr;
+
+    /// the IVF index, controls nlist and nprobe
+    IndexIVF* index_ivf = nullptr;
+
+    /// whether *this owns the 3 fields
+    bool own_fields = false;
+
+    IndexIVFIndependentQuantizer(
+            Index* quantizer,
+            IndexIVF* index_ivf,
+            VectorTransform* vt = nullptr);
+
+    IndexIVFIndependentQuantizer() {}
+
+    void train(idx_t n, const float* x) override;
+
+    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 reset() override;
+
+    ~IndexIVFIndependentQuantizer() override;
+};
+
+} // namespace faiss