faiss 0.2.3 → 0.2.5

data/vendor/faiss/faiss/impl/LocalSearchQuantizer.h

@@ -15,10 +15,17 @@
 #include <vector>
 
 #include <faiss/impl/AdditiveQuantizer.h>
+#include <faiss/impl/platform_macros.h>
 #include <faiss/utils/utils.h>
 
 namespace faiss {
 
+namespace lsq {
+
+struct IcmEncoderFactory;
+
+} // namespace lsq
+
 /** Implementation of LSQ/LSQ++ described in the following two papers:
  *
  * Revisiting additive quantization
@@ -35,7 +42,6 @@ namespace faiss {
  * The trained codes are stored in `codebooks` which is called
  * `centroids` in PQ and RQ.
  */
-
 struct LocalSearchQuantizer : AdditiveQuantizer {
     size_t K; ///< number of codes per codebook
 
@@ -53,10 +59,21 @@ struct LocalSearchQuantizer : AdditiveQuantizer {
     int random_seed; ///< seed for random generator
     size_t nperts;   ///< number of perturbation in each code
 
+    ///< if non-NULL, use this encoder to encode
+    lsq::IcmEncoderFactory* icm_encoder_factory;
+
+    bool update_codebooks_with_double = true;
+
     LocalSearchQuantizer(
-            size_t d,      /* dimensionality of the input vectors */
-            size_t M,      /* number of subquantizers */
-            size_t nbits); /* number of bit per subvector index */
+            size_t d,     /* dimensionality of the input vectors */
+            size_t M,     /* number of subquantizers */
+            size_t nbits, /* number of bit per subvector index */
+            Search_type_t search_type =
+                    ST_decompress); /* determines the storage type */
+
+    LocalSearchQuantizer();
+
+    ~LocalSearchQuantizer() override;
 
     // Train the local search quantizer
     void train(size_t n, const float* x) override;
@@ -65,43 +82,59 @@ struct LocalSearchQuantizer : AdditiveQuantizer {
      *
      * @param x      vectors to encode, size n * d
      * @param codes  output codes, size n * code_size
+     * @param n      number of vectors
+     * @param centroids  centroids to be added to x, size n * d
      */
-    void compute_codes(const float* x, uint8_t* codes, size_t n) const override;
+    void compute_codes_add_centroids(
+            const float* x,
+            uint8_t* codes,
+            size_t n,
+            const float* centroids = nullptr) const override;
 
     /** Update codebooks given encodings
      *
      * @param x      training vectors, size n * d
      * @param codes  encoded training vectors, size n * M
+     * @param n      number of vectors
      */
     void update_codebooks(const float* x, const int32_t* codes, size_t n);
 
     /** Encode vectors given codebooks using iterative conditional mode (icm).
      *
-     * @param x      vectors to encode, size n * d
-     * @param codes  output codes, size n * M
+     * @param codes     output codes, size n * M
+     * @param x         vectors to encode, size n * d
+     * @param n         number of vectors
      * @param ils_iters number of iterations of iterative local search
      */
     void icm_encode(
-            const float* x,
             int32_t* codes,
+            const float* x,
             size_t n,
             size_t ils_iters,
             std::mt19937& gen) const;
 
-    void icm_encode_partial(
-            size_t index,
-            const float* x,
+    void icm_encode_impl(
             int32_t* codes,
+            const float* x,
+            const float* unaries,
+            std::mt19937& gen,
             size_t n,
-            const float* binaries,
             size_t ils_iters,
-            std::mt19937& gen) const;
+            bool verbose) const;
 
     void icm_encode_step(
+            int32_t* codes,
             const float* unaries,
             const float* binaries,
-            int32_t* codes,
-            size_t n) const;
+            size_t n,
+            size_t n_iters) const;
+
+    /** Add some perturbation to codes
+     *
+     * @param codes codes to be perturbed, size n * M
+     * @param n     number of vectors
+     */
+    void perturb_codes(int32_t* codes, size_t n, std::mt19937& gen) const;
 
     /** Add some perturbation to codebooks
      *
@@ -113,12 +146,6 @@ struct LocalSearchQuantizer : AdditiveQuantizer {
             const std::vector<float>& stddev,
             std::mt19937& gen);
 
-    /** Add some perturbation to codes
-     *
-     * @param codes codes to be perturbed, size n * M
-     */
-    void perturb_codes(int32_t* codes, size_t n, std::mt19937& gen) const;
-
     /** Compute binary terms
      *
      * @param binaries binary terms, size M * M * K * K
@@ -127,6 +154,7 @@ struct LocalSearchQuantizer : AdditiveQuantizer {
 
     /** Compute unary terms
      *
+     * @param n       number of vectors
      * @param x       vectors to encode, size n * d
      * @param unaries unary terms, size n * M * K
      */
@@ -134,8 +162,9 @@ struct LocalSearchQuantizer : AdditiveQuantizer {
 
     /** Helper function to compute reconstruction error
      *
-     * @param x     vectors to encode, size n * d
      * @param codes encoded codes, size n * M
+     * @param x     vectors to encode, size n * d
+     * @param n     number of vectors
      * @param objs  if it is not null, store reconstruction
                     error of each vector into it, size n
      */
@@ -146,13 +175,50 @@ struct LocalSearchQuantizer : AdditiveQuantizer {
             float* objs = nullptr) const;
 };
 
+namespace lsq {
+
+struct IcmEncoder {
+    std::vector<float> binaries;
+
+    bool verbose;
+
+    const LocalSearchQuantizer* lsq;
+
+    explicit IcmEncoder(const LocalSearchQuantizer* lsq);
+
+    virtual ~IcmEncoder() {}
+
+    ///< compute binary terms
+    virtual void set_binary_term();
+
+    /** Encode vectors given codebooks
+     *
+     * @param codes     output codes, size n * M
+     * @param x         vectors to encode, size n * d
+     * @param gen       random generator
+     * @param n         number of vectors
+     * @param ils_iters number of iterations of iterative local search
+     */
+    virtual void encode(
+            int32_t* codes,
+            const float* x,
+            std::mt19937& gen,
+            size_t n,
+            size_t ils_iters) const;
+};
+
+struct IcmEncoderFactory {
+    virtual IcmEncoder* get(const LocalSearchQuantizer* lsq) {
+        return new IcmEncoder(lsq);
+    }
+    virtual ~IcmEncoderFactory() {}
+};
+
 /** A helper struct to count consuming time during training.
  *  It is NOT thread-safe.
  */
 struct LSQTimer {
-    std::unordered_map<std::string, double> duration;
-    std::unordered_map<std::string, double> t0;
-    std::unordered_map<std::string, bool> started;
+    std::unordered_map<std::string, double> t;
 
     LSQTimer() {
         reset();
@@ -160,13 +226,24 @@ struct LSQTimer {
 
     double get(const std::string& name);
 
-    void start(const std::string& name);
-
-    void end(const std::string& name);
+    void add(const std::string& name, double delta);
 
     void reset();
 };
 
-FAISS_API extern LSQTimer lsq_timer; ///< timer to count consuming time
+struct LSQTimerScope {
+    double t0;
+    LSQTimer* timer;
+    std::string name;
+    bool finished;
+
+    LSQTimerScope(LSQTimer* timer, std::string name);
+
+    void finish();
+
+    ~LSQTimerScope();
+};
+
+} // namespace lsq
 
 } // namespace faiss

data/vendor/faiss/faiss/impl/LookupTableScaler.h

@@ -0,0 +1,77 @@
+/**
+ * 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 <cstdlib>
+
+#include <faiss/utils/simdlib.h>
+
+/*******************************************
+ * The Scaler objects are used to specialize the handling of the
+ * norm components in Additive quantizer fast-scan.
+ ********************************************/
+
+namespace faiss {
+
+/// no-op handler
+struct DummyScaler {
+    static constexpr int nscale = 0;
+
+    inline simd32uint8 lookup(const simd32uint8&, const simd32uint8&) const {
+        FAISS_THROW_MSG("DummyScaler::lookup should not be called.");
+        return simd32uint8(0);
+    }
+
+    inline simd16uint16 scale_lo(const simd32uint8&) const {
+        FAISS_THROW_MSG("DummyScaler::scale_lo should not be called.");
+        return simd16uint16(0);
+    }
+
+    inline simd16uint16 scale_hi(const simd32uint8&) const {
+        FAISS_THROW_MSG("DummyScaler::scale_hi should not be called.");
+        return simd16uint16(0);
+    }
+
+    template <class dist_t>
+    inline dist_t scale_one(const dist_t&) const {
+        FAISS_THROW_MSG("DummyScaler::scale_one should not be called.");
+        return 0;
+    }
+};
+
+/// consumes 2x4 bits to encode a norm as a scalar additive quantizer
+/// the norm is scaled because its range if larger than other components
+struct NormTableScaler {
+    static constexpr int nscale = 2;
+    int scale_int;
+    simd16uint16 scale_simd;
+
+    explicit NormTableScaler(int scale) : scale_int(scale), scale_simd(scale) {}
+
+    inline simd32uint8 lookup(const simd32uint8& lut, const simd32uint8& c)
+            const {
+        return lut.lookup_2_lanes(c);
+    }
+
+    inline simd16uint16 scale_lo(const simd32uint8& res) const {
+        return simd16uint16(res) * scale_simd;
+    }
+
+    inline simd16uint16 scale_hi(const simd32uint8& res) const {
+        return (simd16uint16(res) >> 8) * scale_simd;
+    }
+
+    // for non-SIMD implem 2, 3, 4
+    template <class dist_t>
+    inline dist_t scale_one(const dist_t& x) const {
+        return x * scale_int;
+    }
+};
+
+} // namespace faiss

data/vendor/faiss/faiss/impl/NNDescent.cpp

@@ -13,6 +13,7 @@
 #include <string>
 
 #include <faiss/impl/AuxIndexStructures.h>
+#include <faiss/impl/DistanceComputer.h>
 
 namespace faiss {
 

data/vendor/faiss/faiss/impl/NSG.cpp

@@ -14,7 +14,7 @@
 #include <mutex>
 #include <stack>
 
-#include <faiss/impl/AuxIndexStructures.h>
+#include <faiss/impl/DistanceComputer.h>
 
 namespace faiss {
 
@@ -161,9 +161,6 @@ void NSG::search(
     search_on_graph<false>(
             *final_graph, dis, vt, enterpoint, pool_size, retset, tmp);
 
-    std::partial_sort(
-            retset.begin(), retset.begin() + k, retset.begin() + pool_size);
-
     for (size_t i = 0; i < k; i++) {
         I[i] = retset[i].id;
         D[i] = retset[i].distance;

data/vendor/faiss/faiss/impl/NSG.h

@@ -105,7 +105,7 @@ struct NSG {
 
     int ntotal; ///< nb of nodes
 
-    /// construction-time parameters
+    // construction-time parameters
     int R; ///< nb of neighbors per node
     int L; ///< length of the search path at construction time
     int C; ///< candidate pool size at construction time