faiss 0.2.6 → 0.2.7

data/vendor/faiss/faiss/impl/AuxIndexStructures.h

@@ -18,7 +18,7 @@
 #include <mutex>
 #include <vector>
 
-#include <faiss/Index.h>
+#include <faiss/MetricType.h>
 #include <faiss/impl/platform_macros.h>
 
 namespace faiss {
@@ -31,15 +31,13 @@ struct RangeSearchResult {
     size_t nq;    ///< nb of queries
     size_t* lims; ///< size (nq + 1)
 
-    typedef Index::idx_t idx_t;
-
     idx_t* labels;    ///< result for query i is labels[lims[i]:lims[i+1]]
     float* distances; ///< corresponding distances (not sorted)
 
     size_t buffer_size; ///< size of the result buffers used
 
     /// lims must be allocated on input to range_search.
-    explicit RangeSearchResult(idx_t nq, bool alloc_lims = true);
+    explicit RangeSearchResult(size_t nq, bool alloc_lims = true);
 
     /// called when lims contains the nb of elements result entries
     /// for each query
@@ -62,8 +60,6 @@ struct RangeSearchResult {
 /** List of temporary buffers used to store results before they are
  *  copied to the RangeSearchResult object. */
 struct BufferList {
-    typedef Index::idx_t idx_t;
-
     // buffer sizes in # entries
     size_t buffer_size;
 
@@ -94,7 +90,6 @@ struct RangeSearchPartialResult;
 
 /// result structure for a single query
 struct RangeQueryResult {
-    using idx_t = Index::idx_t;
     idx_t qno;   //< id of the query
     size_t nres; //< nb of results for this query
     RangeSearchPartialResult* pres;

data/vendor/faiss/faiss/impl/CodePacker.cpp

@@ -0,0 +1,67 @@
+/**
+ * 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/impl/CodePacker.h>
+
+#include <cassert>
+#include <cstring>
+
+namespace faiss {
+
+/*********************************************
+ * CodePacker
+ * default of pack_all / unpack_all loops over the _1 versions
+ */
+
+void CodePacker::pack_all(const uint8_t* flat_codes, uint8_t* block) const {
+    for (size_t i = 0; i < nvec; i++) {
+        pack_1(flat_codes + code_size * i, i, block);
+    }
+}
+
+void CodePacker::unpack_all(const uint8_t* block, uint8_t* flat_codes) const {
+    for (size_t i = 0; i < nvec; i++) {
+        unpack_1(block, i, flat_codes + code_size * i);
+    }
+}
+
+/*********************************************
+ * CodePackerFlat
+ */
+
+CodePackerFlat::CodePackerFlat(size_t code_size) {
+    this->code_size = code_size;
+    nvec = 1;
+    block_size = code_size;
+}
+
+void CodePackerFlat::pack_all(const uint8_t* flat_codes, uint8_t* block) const {
+    memcpy(block, flat_codes, code_size);
+}
+
+void CodePackerFlat::unpack_all(const uint8_t* block, uint8_t* flat_codes)
+        const {
+    memcpy(flat_codes, block, code_size);
+}
+
+void CodePackerFlat::pack_1(
+        const uint8_t* flat_code,
+        size_t offset,
+        uint8_t* block) const {
+    assert(offset == 0);
+    pack_all(flat_code, block);
+}
+
+void CodePackerFlat::unpack_1(
+        const uint8_t* block,
+        size_t offset,
+        uint8_t* flat_code) const {
+    assert(offset == 0);
+    unpack_all(block, flat_code);
+}
+
+} // namespace faiss

data/vendor/faiss/faiss/impl/CodePacker.h

@@ -0,0 +1,71 @@
+/**
+ * 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/MetricType.h>
+
+namespace faiss {
+
+/**
+ * Packing consists in combining a fixed number of codes of constant size
+ * (code_size) into a block of data where they may (or may not) be interleaved
+ * for efficient consumption by distance computation kernels. This exists for
+ * the "fast_scan" indexes on CPU and for some GPU kernels.
+ */
+struct CodePacker {
+    size_t code_size;  // input code size in bytes
+    size_t nvec;       // number of vectors per block
+    size_t block_size; // size of one block in bytes (>= code_size * nvec)
+
+    // pack a single code to a block
+    virtual void pack_1(
+            const uint8_t*
+                    flat_code, // code to write to the block, size code_size
+            size_t offset,     // offset in the block (0 <= offset < nvec)
+            uint8_t* block     // block to write to (size block_size)
+    ) const = 0;
+
+    // unpack a single code from a block
+    virtual void unpack_1(
+            const uint8_t* block, // block to read from (size block_size)
+            size_t offset,        // offset in the block (0 <= offset < nvec)
+            uint8_t* flat_code    // where to write the resulting code, size
+                                  // code_size
+    ) const = 0;
+
+    // pack all code in a block
+    virtual void pack_all(
+            const uint8_t* flat_codes, // codes to write to the block, size
+                                       // (nvec * code_size)
+            uint8_t* block             // block to write to (size block_size)
+    ) const;
+
+    // unpack all code in a block
+    virtual void unpack_all(
+            const uint8_t* block, // block to read from (size block_size)
+            uint8_t* flat_codes // where to write the resulting codes size (nvec
+                                // * code_size)
+    ) const;
+
+    virtual ~CodePacker() {}
+};
+
+/** Trivial code packer where codes are stored one by one */
+struct CodePackerFlat : CodePacker {
+    explicit CodePackerFlat(size_t code_size);
+
+    void pack_1(const uint8_t* flat_code, size_t offset, uint8_t* block)
+            const final;
+    void unpack_1(const uint8_t* block, size_t offset, uint8_t* flat_code)
+            const final;
+
+    void pack_all(const uint8_t* flat_codes, uint8_t* block) const final;
+    void unpack_all(const uint8_t* block, uint8_t* flat_codes) const final;
+};
+
+} // namespace faiss

data/vendor/faiss/faiss/impl/DistanceComputer.h

@@ -23,8 +23,6 @@ namespace faiss {
  * that has additional methods to handle the inverted list context.
  ***********************************************************/
 struct DistanceComputer {
-    using idx_t = Index::idx_t;
-
     /// called before computing distances. Pointer x should remain valid
     /// while operator () is called
     virtual void set_query(const float* x) = 0;

data/vendor/faiss/faiss/impl/HNSW.cpp

@@ -47,11 +47,6 @@ void HNSW::neighbor_range(idx_t no, int layer_no, size_t* begin, size_t* end)
 
 HNSW::HNSW(int M) : rng(12345) {
     set_default_probas(M, 1.0 / log(M));
-    max_level = -1;
-    entry_point = -1;
-    efSearch = 16;
-    efConstruction = 40;
-    upper_beam = 1;
     offsets.push_back(0);
 }
 
@@ -509,7 +504,6 @@ void HNSW::add_with_locks(
 
 namespace {
 
-using idx_t = HNSW::idx_t;
 using MinimaxHeap = HNSW::MinimaxHeap;
 using Node = HNSW::Node;
 /** Do a BFS on the candidates list */
@@ -837,8 +831,10 @@ void HNSW::MinimaxHeap::push(storage_idx_t i, float v) {
     if (k == n) {
         if (v >= dis[0])
             return;
+        if (ids[0] != -1) {
+            --nvalid;
+        }
         faiss::heap_pop<HC>(k--, dis.data(), ids.data());
-        --nvalid;
     }
     faiss::heap_push<HC>(++k, dis.data(), ids.data(), v, i);
     ++nvalid;

data/vendor/faiss/faiss/impl/HNSW.h

@@ -52,10 +52,7 @@ struct SearchParametersHNSW : SearchParameters {
 
 struct HNSW {
     /// internal storage of vectors (32 bits: this is expensive)
-    typedef int storage_idx_t;
-
-    /// Faiss results are 64-bit
-    typedef Index::idx_t idx_t;
+    using storage_idx_t = int32_t;
 
     typedef std::pair<float, storage_idx_t> Node;
 
@@ -124,25 +121,25 @@ struct HNSW {
 
     /// entry point in the search structure (one of the points with maximum
     /// level
-    storage_idx_t entry_point;
+    storage_idx_t entry_point = -1;
 
     faiss::RandomGenerator rng;
 
     /// maximum level
-    int max_level;
+    int max_level = -1;
 
     /// expansion factor at construction time
-    int efConstruction;
+    int efConstruction = 40;
 
     /// expansion factor at search time
-    int efSearch;
+    int efSearch = 16;
 
     /// during search: do we check whether the next best distance is good
     /// enough?
     bool check_relative_distance = true;
 
     /// number of entry points in levels > 0.
-    int upper_beam;
+    int upper_beam = 1;
 
     /// use bounded queue during exploration
     bool search_bounded_queue = true;

data/vendor/faiss/faiss/impl/IDSelector.cpp

@@ -92,7 +92,7 @@ IDSelectorBatch::IDSelectorBatch(size_t n, const idx_t* indices) {
     mask = ((idx_t)1 << nbits) - 1;
     bloom.resize((idx_t)1 << (nbits - 3), 0);
     for (idx_t i = 0; i < n; i++) {
-        Index::idx_t id = indices[i];
+        idx_t id = indices[i];
         set.insert(id);
         id &= mask;
         bloom[id >> 3] |= 1 << (id & 7);

data/vendor/faiss/faiss/impl/IDSelector.h

@@ -19,7 +19,6 @@ namespace faiss {
 
 /** Encapsulates a set of ids to handle. */
 struct IDSelector {
-    using idx_t = Index::idx_t;
     virtual bool is_member(idx_t id) const = 0;
     virtual ~IDSelector() {}
 };
@@ -132,4 +131,43 @@ struct IDSelectorAll : IDSelector {
     virtual ~IDSelectorAll() {}
 };
 
+/// does an AND operation on the the two given IDSelector's is_membership
+/// results.
+struct IDSelectorAnd : IDSelector {
+    const IDSelector* lhs;
+    const IDSelector* rhs;
+    IDSelectorAnd(const IDSelector* lhs, const IDSelector* rhs)
+            : lhs(lhs), rhs(rhs) {}
+    bool is_member(idx_t id) const final {
+        return lhs->is_member(id) && rhs->is_member(id);
+    };
+    virtual ~IDSelectorAnd() {}
+};
+
+/// does an OR operation on the the two given IDSelector's is_membership
+/// results.
+struct IDSelectorOr : IDSelector {
+    const IDSelector* lhs;
+    const IDSelector* rhs;
+    IDSelectorOr(const IDSelector* lhs, const IDSelector* rhs)
+            : lhs(lhs), rhs(rhs) {}
+    bool is_member(idx_t id) const final {
+        return lhs->is_member(id) || rhs->is_member(id);
+    };
+    virtual ~IDSelectorOr() {}
+};
+
+/// does an XOR operation on the the two given IDSelector's is_membership
+/// results.
+struct IDSelectorXOr : IDSelector {
+    const IDSelector* lhs;
+    const IDSelector* rhs;
+    IDSelectorXOr(const IDSelector* lhs, const IDSelector* rhs)
+            : lhs(lhs), rhs(rhs) {}
+    bool is_member(idx_t id) const final {
+        return lhs->is_member(id) ^ rhs->is_member(id);
+    };
+    virtual ~IDSelectorXOr() {}
+};
+
 } // namespace faiss

data/vendor/faiss/faiss/impl/LocalSearchQuantizer.cpp

@@ -21,6 +21,15 @@
 #include <faiss/utils/hamming.h> // BitstringWriter
 #include <faiss/utils/utils.h>
 
+#include <faiss/utils/approx_topk/approx_topk.h>
+
+// this is needed for prefetching
+#include <faiss/impl/platform_macros.h>
+
+#ifdef __AVX2__
+#include <xmmintrin.h>
+#endif
+
 extern "C" {
 // LU decomoposition of a general matrix
 void sgetrf_(
@@ -151,23 +160,7 @@ LocalSearchQuantizer::LocalSearchQuantizer(
         Search_type_t search_type)
         : AdditiveQuantizer(d, std::vector<size_t>(M, nbits), search_type) {
     K = (1 << nbits);
-
-    train_iters = 25;
-    train_ils_iters = 8;
-    icm_iters = 4;
-
-    encode_ils_iters = 16;
-
-    p = 0.5f;
-    lambd = 1e-2f;
-
-    chunk_size = 10000;
-    nperts = 4;
-
-    random_seed = 0x12345;
     std::srand(random_seed);
-
-    icm_encoder_factory = nullptr;
 }
 
 LocalSearchQuantizer::~LocalSearchQuantizer() {
@@ -192,7 +185,7 @@ void LocalSearchQuantizer::train(size_t n, const float* x) {
     // allocate memory for codebooks, size [M, K, d]
     codebooks.resize(M * K * d);
 
-    // randomly intialize codes
+    // randomly initialize codes
     std::mt19937 gen(random_seed);
     std::vector<int32_t> codes(n * M); // [n, M]
     random_int32(codes, 0, K - 1, gen);
@@ -604,54 +597,72 @@ void LocalSearchQuantizer::icm_encode_step(
     FAISS_THROW_IF_NOT(M != 0 && K != 0);
     FAISS_THROW_IF_NOT(binaries != nullptr);
 
-    for (size_t iter = 0; iter < n_iters; iter++) {
-        // condition on the m-th subcode
-        for (size_t m = 0; m < M; m++) {
-            std::vector<float> objs(n * K);
-#pragma omp parallel for
-            for (int64_t i = 0; i < n; i++) {
-                auto u = unaries + m * n * K + i * K;
-                memcpy(objs.data() + i * K, u, sizeof(float) * K);
-            }
+#pragma omp parallel for schedule(dynamic)
+    for (int64_t i = 0; i < n; i++) {
+        std::vector<float> objs(K);
 
-            // compute objective function by adding unary
-            // and binary terms together
-            for (size_t other_m = 0; other_m < M; other_m++) {
-                if (other_m == m) {
-                    continue;
+        for (size_t iter = 0; iter < n_iters; iter++) {
+            // condition on the m-th subcode
+            for (size_t m = 0; m < M; m++) {
+                // copy
+                auto u = unaries + m * n * K + i * K;
+                for (size_t code = 0; code < K; code++) {
+                    objs[code] = u[code];
                 }
 
-#pragma omp parallel for
-                for (int64_t i = 0; i < n; i++) {
+                // compute objective function by adding unary
+                // and binary terms together
+                for (size_t other_m = 0; other_m < M; other_m++) {
+                    if (other_m == m) {
+                        continue;
+                    }
+
+#ifdef __AVX2__
+                    // TODO: add platform-independent compiler-independent
+                    // prefetch utilities.
+                    if (other_m + 1 < M) {
+                        // do a single prefetch
+                        int32_t code2 = codes[i * M + other_m + 1];
+                        // for (int32_t code = 0; code < K; code += 64) {
+                        int32_t code = 0;
+                        {
+                            size_t binary_idx = (other_m + 1) * M * K * K +
+                                    m * K * K + code2 * K + code;
+                            _mm_prefetch(binaries + binary_idx, _MM_HINT_T0);
+                        }
+                    }
+#endif
+
                     for (int32_t code = 0; code < K; code++) {
                         int32_t code2 = codes[i * M + other_m];
-                        size_t binary_idx = m * M * K * K + other_m * K * K +
-                                code * K + code2;
-                        // binaries[m, other_m, code, code2]
-                        objs[i * K + code] += binaries[binary_idx];
+                        size_t binary_idx = other_m * M * K * K + m * K * K +
+                                code2 * K + code;
+                        // binaries[m, other_m, code, code2].
+                        // It is symmetric over (m <-> other_m)
+                        //   and (code <-> code2).
+                        // So, replace the op with
+                        //   binaries[other_m, m, code2, code].
+                        objs[code] += binaries[binary_idx];
                     }
                 }
-            }
 
-            // find the optimal value of the m-th subcode
-#pragma omp parallel for
-            for (int64_t i = 0; i < n; i++) {
+                // find the optimal value of the m-th subcode
                 float best_obj = HUGE_VALF;
                 int32_t best_code = 0;
-                for (size_t code = 0; code < K; code++) {
-                    float obj = objs[i * K + code];
-                    if (obj < best_obj) {
-                        best_obj = obj;
-                        best_code = code;
-                    }
-                }
+
+                // find one using SIMD. The following operation is similar
+                // to the search of the smallest element in objs
+                using C = CMax<float, int>;
+                HeapWithBuckets<C, 16, 1>::addn(
+                        K, objs.data(), 1, &best_obj, &best_code);
+
+                // done
                 codes[i * M + m] = best_code;
-            }
 
-        } // loop M
+            } // loop M
+        }
     }
 }
-
 void LocalSearchQuantizer::perturb_codes(
         int32_t* codes,
         size_t n,

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

@@ -45,22 +45,21 @@ struct IcmEncoderFactory;
 struct LocalSearchQuantizer : AdditiveQuantizer {
     size_t K; ///< number of codes per codebook
 
-    size_t train_iters; ///< number of iterations in training
+    size_t train_iters = 25;      ///< number of iterations in training
+    size_t encode_ils_iters = 16; ///< iterations of local search in encoding
+    size_t train_ils_iters = 8;   ///< iterations of local search in training
+    size_t icm_iters = 4;         ///< number of iterations in icm
 
-    size_t encode_ils_iters; ///< iterations of local search in encoding
-    size_t train_ils_iters;  ///< iterations of local search in training
-    size_t icm_iters;        ///< number of iterations in icm
+    float p = 0.5f;      ///< temperature factor
+    float lambd = 1e-2f; ///< regularization factor
 
-    float p;     ///< temperature factor
-    float lambd; ///< regularization factor
+    size_t chunk_size = 10000; ///< nb of vectors to encode at a time
 
-    size_t chunk_size; ///< nb of vectors to encode at a time
+    int random_seed = 0x12345; ///< seed for random generator
+    size_t nperts = 4;         ///< number of perturbation in each code
 
-    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;
+    ///< if non-NULL, use this encoder to encode (owned by the object)
+    lsq::IcmEncoderFactory* icm_encoder_factory = nullptr;
 
     bool update_codebooks_with_double = true;
 

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

@@ -147,14 +147,8 @@ using namespace nndescent;
 
 constexpr int NUM_EVAL_POINTS = 100;
 
-NNDescent::NNDescent(const int d, const int K) : K(K), random_seed(2021), d(d) {
-    ntotal = 0;
-    has_built = false;
-    S = 10;
-    R = 100;
+NNDescent::NNDescent(const int d, const int K) : K(K), d(d) {
     L = K + 50;
-    iter = 10;
-    search_L = 0;
 }
 
 NNDescent::~NNDescent() {}
@@ -311,7 +305,7 @@ void NNDescent::generate_eval_set(
     for (int i = 0; i < c.size(); i++) {
         std::vector<Neighbor> tmp;
         for (int j = 0; j < N; j++) {
-            if (i == j)
+            if (c[i] == j)
                 continue; // skip itself
             float dist = qdis.symmetric_dis(c[i], j);
             tmp.push_back(Neighbor(j, dist, true));
@@ -425,7 +419,7 @@ void NNDescent::search(
     // candidate pool, the K best items is the result.
     std::vector<Neighbor> retset(L + 1);
 
-    // Randomly choose L points to intialize the candidate pool
+    // Randomly choose L points to initialize the candidate pool
     std::vector<int> init_ids(L);
     std::mt19937 rng(random_seed);
 

data/vendor/faiss/faiss/impl/NNDescent.h

@@ -90,7 +90,6 @@ struct Nhood {
 
 struct NNDescent {
     using storage_idx_t = int;
-    using idx_t = Index::idx_t;
 
     using KNNGraph = std::vector<nndescent::Nhood>;
 
@@ -133,19 +132,20 @@ struct NNDescent {
             std::vector<int>& ctrl_points,
             std::vector<std::vector<int>>& acc_eval_set);
 
-    bool has_built;
+    bool has_built = false;
 
-    int K; // K in KNN graph
-    int S; // number of sample neighbors to be updated for each node
-    int R; // size of reverse links, 0 means the reverse links will not be used
-    int L; // size of the candidate pool in building
-    int iter;        // number of iterations to iterate over
-    int search_L;    // size of candidate pool in searching
-    int random_seed; // random seed for generators
+    int S = 10;  // number of sample neighbors to be updated for each node
+    int R = 100; // size of reverse links, 0 means the reverse links will not be
+                 // used
+    int iter = 10;          // number of iterations to iterate over
+    int search_L = 0;       // size of candidate pool in searching
+    int random_seed = 2021; // random seed for generators
 
+    int K; // K in KNN graph
     int d; // dimensions
+    int L; // size of the candidate pool in building
 
-    int ntotal;
+    int ntotal = 0;
 
     KNNGraph graph;
     std::vector<int> final_graph;

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

@@ -29,8 +29,6 @@ constexpr int EMPTY_ID = -1;
    distances. This makes supporting INNER_PRODUCE search easier */
 
 struct NegativeDistanceComputer : DistanceComputer {
-    using idx_t = Index::idx_t;
-
     /// owned by this
     DistanceComputer* basedis;
 
@@ -59,7 +57,7 @@ struct NegativeDistanceComputer : DistanceComputer {
 } // namespace
 
 DistanceComputer* storage_distance_computer(const Index* storage) {
-    if (storage->metric_type == METRIC_INNER_PRODUCT) {
+    if (is_similarity_metric(storage->metric_type)) {
         return new NegativeDistanceComputer(storage->get_distance_computer());
     } else {
         return storage->get_distance_computer();
@@ -140,9 +138,6 @@ inline int insert_into_pool(Neighbor* addr, int K, Neighbor nn) {
 NSG::NSG(int R) : R(R), rng(0x0903) {
     L = R + 32;
     C = R + 100;
-    search_L = 16;
-    ntotal = 0;
-    is_built = false;
     srand(0x1998);
 }
 

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

@@ -98,12 +98,9 @@ DistanceComputer* storage_distance_computer(const Index* storage);
 
 struct NSG {
     /// internal storage of vectors (32 bits: this is expensive)
-    using storage_idx_t = int;
+    using storage_idx_t = int32_t;
 
-    /// Faiss results are 64-bit
-    using idx_t = Index::idx_t;
-
-    int ntotal; ///< nb of nodes
+    int ntotal = 0; ///< nb of nodes
 
     // construction-time parameters
     int R; ///< nb of neighbors per node
@@ -111,13 +108,13 @@ struct NSG {
     int C; ///< candidate pool size at construction time
 
     // search-time parameters
-    int search_L; ///< length of the search path
+    int search_L = 16; ///< length of the search path
 
     int enterpoint; ///< enterpoint
 
     std::shared_ptr<nsg::Graph<int>> final_graph; ///< NSG graph structure
 
-    bool is_built; ///< NSG is built or not
+    bool is_built = false; ///< NSG is built or not
 
     RandomGenerator rng; ///< random generator