faiss 0.2.0 → 0.2.1

data/vendor/faiss/faiss/impl/ResidualQuantizer.h

@@ -0,0 +1,130 @@
+/**
+ * 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 <cstdint>
+#include <vector>
+
+#include <faiss/Clustering.h>
+#include <faiss/impl/AdditiveQuantizer.h>
+
+namespace faiss {
+
+/** Residual quantizer with variable number of bits per sub-quantizer
+ *
+ * The residual centroids are stored in a big cumulative centroid table.
+ * The codes are represented either as a non-compact table of size (n, M) or
+ * as the compact output (n, code_size).
+ */
+
+struct ResidualQuantizer : AdditiveQuantizer {
+    /// initialization
+    enum train_type_t {
+        Train_default,         ///< regular k-means
+        Train_progressive_dim, ///< progressive dim clustering
+    };
+
+    // set this bit on train_type if beam is to be trained only on the
+    // first element of the beam (faster but less accurate)
+    static const int Train_top_beam = 1024;
+    train_type_t train_type;
+
+    /// beam size used for training and for encoding
+    int max_beam_size;
+
+    /// distance matrixes with beam search can get large, so use this
+    /// to batch computations at encoding time.
+    size_t max_mem_distances;
+
+    /// clustering parameters
+    ProgressiveDimClusteringParameters cp;
+
+    /// if non-NULL, use this index for assignment
+    ProgressiveDimIndexFactory* assign_index_factory;
+
+    ResidualQuantizer(size_t d, const std::vector<size_t>& nbits);
+
+    ResidualQuantizer(
+            size_t d,      /* dimensionality of the input vectors */
+            size_t M,      /* number of subquantizers */
+            size_t nbits); /* number of bit per subvector index */
+
+    ResidualQuantizer();
+
+    // Train the residual quantizer
+    void train(size_t n, const float* x) override;
+
+    /** Encode a set of vectors
+     *
+     * @param x      vectors to encode, size n * d
+     * @param codes  output codes, size n * code_size
+     */
+    void compute_codes(const float* x, uint8_t* codes, size_t n) const override;
+
+    /** lower-level encode function
+     *
+     * @param n              number of vectors to hanlde
+     * @param residuals      vectors to encode, size (n, beam_size, d)
+     * @param beam_size      input beam size
+     * @param new_beam_size  output beam size (should be <= K * beam_size)
+     * @param new_codes      output codes, size (n, new_beam_size, m + 1)
+     * @param new_residuals  output residuals, size (n, new_beam_size, d)
+     * @param new_distances  output distances, size (n, new_beam_size)
+     */
+    void refine_beam(
+            size_t n,
+            size_t beam_size,
+            const float* residuals,
+            int new_beam_size,
+            int32_t* new_codes,
+            float* new_residuals = nullptr,
+            float* new_distances = nullptr) const;
+
+    /** Beam search can consume a lot of memory. This function estimates the
+     * amount of mem used by refine_beam to adjust the batch size
+     *
+     * @param beam_size  if != -1, override the beam size
+     */
+    size_t memory_per_point(int beam_size = -1) const;
+};
+
+/** Encode a residual by sampling from a centroid table.
+ *
+ * This is a single encoding step the residual quantizer.
+ * It allows low-level access to the encoding function, exposed mainly for unit
+ * tests.
+ *
+ * @param n              number of vectors to hanlde
+ * @param residuals      vectors to encode, size (n, beam_size, d)
+ * @param cent           centroids, size (K, d)
+ * @param beam_size      input beam size
+ * @param m              size of the codes for the previous encoding steps
+ * @param codes          code array for the previous steps of the beam (n,
+ * beam_size, m)
+ * @param new_beam_size  output beam size (should be <= K * beam_size)
+ * @param new_codes      output codes, size (n, new_beam_size, m + 1)
+ * @param new_residuals  output residuals, size (n, new_beam_size, d)
+ * @param new_distances  output distances, size (n, new_beam_size)
+ * @param assign_index   if non-NULL, will be used to perform assignment
+ */
+void beam_search_encode_step(
+        size_t d,
+        size_t K,
+        const float* cent,
+        size_t n,
+        size_t beam_size,
+        const float* residuals,
+        size_t m,
+        const int32_t* codes,
+        size_t new_beam_size,
+        int32_t* new_codes,
+        float* new_residuals,
+        float* new_distances,
+        Index* assign_index = nullptr);
+
+}; // namespace faiss

data/vendor/faiss/faiss/impl/ResultHandler.h

@@ -5,49 +5,38 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 /*
  * Structures that collect search results from distance computations
  */
 
 #pragma once
 
-
+#include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/utils/Heap.h>
 #include <faiss/utils/partitioning.h>
-#include <faiss/impl/AuxIndexStructures.h>
-
 
 namespace faiss {
 
-
-
 /*****************************************************************
  * Heap based result handler
  *****************************************************************/
 
-
-template<class C>
+template <class C>
 struct HeapResultHandler {
-
     using T = typename C::T;
     using TI = typename C::TI;
 
     int nq;
-    T *heap_dis_tab;
-    TI *heap_ids_tab;
+    T* heap_dis_tab;
+    TI* heap_ids_tab;
 
-    int64_t k;  // number of results to keep
+    int64_t k; // number of results to keep
 
-    HeapResultHandler(
-        size_t nq,
-        T * heap_dis_tab, TI * heap_ids_tab,
-        size_t k):
-        nq(nq),
-        heap_dis_tab(heap_dis_tab), heap_ids_tab(heap_ids_tab), k(k)
-    {
-
-    }
+    HeapResultHandler(size_t nq, T* heap_dis_tab, TI* heap_ids_tab, size_t k)
+            : nq(nq),
+              heap_dis_tab(heap_dis_tab),
+              heap_ids_tab(heap_ids_tab),
+              k(k) {}
 
     /******************************************************
      * API for 1 result at a time (each SingleResultHandler is
@@ -55,20 +44,20 @@ struct HeapResultHandler {
      */
 
     struct SingleResultHandler {
-        HeapResultHandler & hr;
+        HeapResultHandler& hr;
         size_t k;
 
-        T *heap_dis;
-        TI *heap_ids;
+        T* heap_dis;
+        TI* heap_ids;
         T thresh;
 
-        SingleResultHandler(HeapResultHandler &hr): hr(hr), k(hr.k) {}
+        SingleResultHandler(HeapResultHandler& hr) : hr(hr), k(hr.k) {}
 
         /// begin results for query # i
         void begin(size_t i) {
             heap_dis = hr.heap_dis_tab + i * k;
             heap_ids = hr.heap_ids_tab + i * k;
-            heap_heapify<C> (k, heap_dis, heap_ids);
+            heap_heapify<C>(k, heap_dis, heap_ids);
             thresh = heap_dis[0];
         }
 
@@ -82,11 +71,10 @@ struct HeapResultHandler {
 
         /// series of results for query i is done
         void end() {
-            heap_reorder<C> (k, heap_dis, heap_ids);
+            heap_reorder<C>(k, heap_dis, heap_ids);
         }
     };
 
-
     /******************************************************
      * API for multiple results (called from 1 thread)
      */
@@ -97,20 +85,21 @@ struct HeapResultHandler {
     void begin_multiple(size_t i0, size_t i1) {
         this->i0 = i0;
         this->i1 = i1;
-        for(size_t i = i0; i < i1; i++) {
-            heap_heapify<C> (k, heap_dis_tab + i * k, heap_ids_tab + i * k);
+        for (size_t i = i0; i < i1; i++) {
+            heap_heapify<C>(k, heap_dis_tab + i * k, heap_ids_tab + i * k);
         }
     }
 
     /// add results for query i0..i1 and j0..j1
-    void add_results(size_t j0, size_t j1, const T *dis_tab) {
-        // maybe parallel for
-        for (size_t i = i0; i < i1; i++) {
-            T * heap_dis = heap_dis_tab + i * k;
-            TI * heap_ids = heap_ids_tab + i * k;
+    void add_results(size_t j0, size_t j1, const T* dis_tab) {
+#pragma omp parallel for
+        for (int64_t i = i0; i < i1; i++) {
+            T* heap_dis = heap_dis_tab + i * k;
+            TI* heap_ids = heap_ids_tab + i * k;
+            const T* dis_tab_i = dis_tab + (j1 - j0) * (i - i0) - j0;
             T thresh = heap_dis[0];
             for (size_t j = j0; j < j1; j++) {
-                T dis = *dis_tab++;
+                T dis = dis_tab_i[j];
                 if (C::cmp(thresh, dis)) {
                     heap_replace_top<C>(k, heap_dis, heap_ids, dis, j);
                     thresh = heap_dis[0];
@@ -122,11 +111,10 @@ struct HeapResultHandler {
     /// series of results for queries i0..i1 is done
     void end_multiple() {
         // maybe parallel for
-        for(size_t i = i0; i < i1; i++) {
-            heap_reorder<C> (k, heap_dis_tab + i * k, heap_ids_tab + i * k);
+        for (size_t i = i0; i < i1; i++) {
+            heap_reorder<C>(k, heap_dis_tab + i * k, heap_ids_tab + i * k);
         }
     }
-
 };
 
 /*****************************************************************
@@ -138,31 +126,25 @@ struct HeapResultHandler {
  * distance array.
  *****************************************************************/
 
-
-
 /// Reservoir for a single query
-template<class C>
+template <class C>
 struct ReservoirTopN {
     using T = typename C::T;
     using TI = typename C::TI;
 
-    T *vals;
-    TI *ids;
+    T* vals;
+    TI* ids;
 
-    size_t i; // number of stored elements
-    size_t n; // number of requested elements
-    size_t capacity;  // size of storage
+    size_t i;        // number of stored elements
+    size_t n;        // number of requested elements
+    size_t capacity; // size of storage
 
     T threshold; // current threshold
 
     ReservoirTopN() {}
 
-    ReservoirTopN(
-        size_t n, size_t capacity,
-        T *vals, TI *ids
-    ):
-        vals(vals), ids(ids),
-        i(0), n(n), capacity(capacity) {
+    ReservoirTopN(size_t n, size_t capacity, T* vals, TI* ids)
+            : vals(vals), ids(ids), i(0), n(n), capacity(capacity) {
         assert(n < capacity);
         threshold = C::neutral();
     }
@@ -184,55 +166,47 @@ struct ReservoirTopN {
         assert(i == capacity);
 
         threshold = partition_fuzzy<C>(
-            vals, ids, capacity, n, (capacity + n) / 2,
-            &i);
+                vals, ids, capacity, n, (capacity + n) / 2, &i);
     }
 
-    void to_result(T *heap_dis, TI *heap_ids) const {
-
+    void to_result(T* heap_dis, TI* heap_ids) const {
         for (int j = 0; j < std::min(i, n); j++) {
-            heap_push<C>(
-                j + 1, heap_dis, heap_ids,
-                vals[j], ids[j]
-            );
+            heap_push<C>(j + 1, heap_dis, heap_ids, vals[j], ids[j]);
         }
 
         if (i < n) {
-            heap_reorder<C> (i, heap_dis, heap_ids);
+            heap_reorder<C>(i, heap_dis, heap_ids);
             // add empty results
-            heap_heapify<C> (n - i, heap_dis + i, heap_ids + i);
+            heap_heapify<C>(n - i, heap_dis + i, heap_ids + i);
         } else {
             // add remaining elements
-            heap_addn<C> (n, heap_dis, heap_ids, vals + n, ids + n, i - n);
-            heap_reorder<C> (n, heap_dis, heap_ids);
+            heap_addn<C>(n, heap_dis, heap_ids, vals + n, ids + n, i - n);
+            heap_reorder<C>(n, heap_dis, heap_ids);
         }
-
     }
-
 };
 
-
-
-template<class C>
+template <class C>
 struct ReservoirResultHandler {
-
     using T = typename C::T;
     using TI = typename C::TI;
 
     int nq;
-    T *heap_dis_tab;
-    TI *heap_ids_tab;
+    T* heap_dis_tab;
+    TI* heap_ids_tab;
 
-    int64_t k;  // number of results to keep
+    int64_t k;       // number of results to keep
     size_t capacity; // capacity of the reservoirs
 
     ReservoirResultHandler(
-        size_t nq,
-        T * heap_dis_tab, TI * heap_ids_tab,
-        size_t k):
-        nq(nq),
-        heap_dis_tab(heap_dis_tab), heap_ids_tab(heap_ids_tab), k(k)
-    {
+            size_t nq,
+            T* heap_dis_tab,
+            TI* heap_ids_tab,
+            size_t k)
+            : nq(nq),
+              heap_dis_tab(heap_dis_tab),
+              heap_ids_tab(heap_ids_tab),
+              k(k) {
         // double then round up to multiple of 16 (for SIMD alignment)
         capacity = (2 * k + 15) & ~15;
     }
@@ -243,23 +217,26 @@ struct ReservoirResultHandler {
      */
 
     struct SingleResultHandler {
-        ReservoirResultHandler & hr;
+        ReservoirResultHandler& hr;
 
         std::vector<T> reservoir_dis;
         std::vector<TI> reservoir_ids;
         ReservoirTopN<C> res1;
 
-        SingleResultHandler(ReservoirResultHandler &hr):
-            hr(hr), reservoir_dis(hr.capacity), reservoir_ids(hr.capacity)
-        {
-        }
+        SingleResultHandler(ReservoirResultHandler& hr)
+                : hr(hr),
+                  reservoir_dis(hr.capacity),
+                  reservoir_ids(hr.capacity) {}
 
         size_t i;
 
         /// begin results for query # i
         void begin(size_t i) {
             res1 = ReservoirTopN<C>(
-                hr.k, hr.capacity, reservoir_dis.data(), reservoir_ids.data());
+                    hr.k,
+                    hr.capacity,
+                    reservoir_dis.data(),
+                    reservoir_ids.data());
             this->i = i;
         }
 
@@ -270,8 +247,8 @@ struct ReservoirResultHandler {
 
         /// series of results for query i is done
         void end() {
-            T * heap_dis = hr.heap_dis_tab + i * hr.k;
-            TI * heap_ids = hr.heap_ids_tab + i * hr.k;
+            T* heap_dis = hr.heap_dis_tab + i * hr.k;
+            TI* heap_ids = hr.heap_ids_tab + i * hr.k;
             res1.to_result(heap_dis, heap_ids);
         }
     };
@@ -295,20 +272,22 @@ struct ReservoirResultHandler {
         reservoirs.clear();
         for (size_t i = i0; i < i1; i++) {
             reservoirs.emplace_back(
-                k, capacity,
-                reservoir_dis.data() + (i - i0) * capacity,
-                reservoir_ids.data() + (i - i0) * capacity
-            );
+                    k,
+                    capacity,
+                    reservoir_dis.data() + (i - i0) * capacity,
+                    reservoir_ids.data() + (i - i0) * capacity);
         }
     }
 
     /// add results for query i0..i1 and j0..j1
-    void add_results(size_t j0, size_t j1, const T *dis_tab) {
+    void add_results(size_t j0, size_t j1, const T* dis_tab) {
         // maybe parallel for
-        for (size_t i = i0; i < i1; i++) {
-            ReservoirTopN<C> & reservoir = reservoirs[i - i0];
+#pragma omp parallel for
+        for (int64_t i = i0; i < i1; i++) {
+            ReservoirTopN<C>& reservoir = reservoirs[i - i0];
+            const T* dis_tab_i = dis_tab + (j1 - j0) * (i - i0) - j0;
             for (size_t j = j0; j < j1; j++) {
-                T dis = *dis_tab++;
+                T dis = dis_tab_i[j];
                 reservoir.add(dis, j);
             }
         }
@@ -317,32 +296,27 @@ struct ReservoirResultHandler {
     /// series of results for queries i0..i1 is done
     void end_multiple() {
         // maybe parallel for
-        for(size_t i = i0; i < i1; i++) {
+        for (size_t i = i0; i < i1; i++) {
             reservoirs[i - i0].to_result(
-                heap_dis_tab + i * k, heap_ids_tab + i * k);
+                    heap_dis_tab + i * k, heap_ids_tab + i * k);
         }
     }
-
 };
 
-
 /*****************************************************************
  * Result handler for range searches
  *****************************************************************/
 
-
-
-template<class C>
+template <class C>
 struct RangeSearchResultHandler {
     using T = typename C::T;
     using TI = typename C::TI;
 
-    RangeSearchResult *res;
+    RangeSearchResult* res;
     float radius;
 
-    RangeSearchResultHandler(RangeSearchResult *res, float radius):
-        res(res), radius(radius)
-    {}
+    RangeSearchResultHandler(RangeSearchResult* res, float radius)
+            : res(res), radius(radius) {}
 
     /******************************************************
      * API for 1 result at a time (each SingleResultHandler is
@@ -353,11 +327,10 @@ struct RangeSearchResultHandler {
         // almost the same interface as RangeSearchResultHandler
         RangeSearchPartialResult pres;
         float radius;
-        RangeQueryResult *qr = nullptr;
+        RangeQueryResult* qr = nullptr;
 
-        SingleResultHandler(RangeSearchResultHandler &rh):
-            pres(rh.res), radius(rh.radius)
-        {}
+        SingleResultHandler(RangeSearchResultHandler& rh)
+                : pres(rh.res), radius(rh.radius) {}
 
         /// begin results for query # i
         void begin(size_t i) {
@@ -366,15 +339,13 @@ struct RangeSearchResultHandler {
 
         /// add one result for query i
         void add_result(T dis, TI idx) {
-
             if (C::cmp(radius, dis)) {
                 qr->add(dis, idx);
             }
         }
 
         /// series of results for query i is done
-        void end() {
-        }
+        void end() {}
 
         ~SingleResultHandler() {
             pres.finalize();
@@ -387,8 +358,8 @@ struct RangeSearchResultHandler {
 
     size_t i0, i1;
 
-    std::vector <RangeSearchPartialResult *> partial_results;
-    std::vector <size_t> j0s;
+    std::vector<RangeSearchPartialResult*> partial_results;
+    std::vector<size_t> j0s;
     int pr = 0;
 
     /// begin
@@ -399,8 +370,8 @@ struct RangeSearchResultHandler {
 
     /// add results for query i0..i1 and j0..j1
 
-    void add_results(size_t j0, size_t j1, const T *dis_tab) {
-        RangeSearchPartialResult *pres;
+    void add_results(size_t j0, size_t j1, const T* dis_tab) {
+        RangeSearchPartialResult* pres;
         // there is one RangeSearchPartialResult structure per j0
         // (= block of columns of the large distance matrix)
         // it is a bit tricky to find the poper PartialResult structure
@@ -414,39 +385,32 @@ struct RangeSearchResultHandler {
             pres = partial_results[pr];
             pr++;
         } else { // did not find this j0
-            pres = new RangeSearchPartialResult (res);
+            pres = new RangeSearchPartialResult(res);
             partial_results.push_back(pres);
             j0s.push_back(j0);
             pr = partial_results.size();
         }
 
         for (size_t i = i0; i < i1; i++) {
-            const float *ip_line = dis_tab + (i - i0) * (j1 - j0);
-            RangeQueryResult & qres = pres->new_result (i);
+            const float* ip_line = dis_tab + (i - i0) * (j1 - j0);
+            RangeQueryResult& qres = pres->new_result(i);
 
             for (size_t j = j0; j < j1; j++) {
                 float dis = *ip_line++;
                 if (C::cmp(radius, dis)) {
-                    qres.add (dis, j);
+                    qres.add(dis, j);
                 }
             }
         }
     }
 
-    void end_multiple() {
-
-    }
+    void end_multiple() {}
 
     ~RangeSearchResultHandler() {
         if (partial_results.size() > 0) {
-            RangeSearchPartialResult::merge (partial_results);
+            RangeSearchPartialResult::merge(partial_results);
         }
     }
-
 };
 
-
-
-
-}  // namespace faiss
-
+} // namespace faiss