faiss 0.3.0 → 0.3.2

data/vendor/faiss/faiss/IndexNSG.h

@@ -24,29 +24,29 @@ namespace faiss {
  * link structure built on top */
 
 struct IndexNSG : Index {
-    /// the link strcuture
+    /// the link structure
     NSG nsg;
 
     /// the sequential storage
-    bool own_fields;
-    Index* storage;
+    bool own_fields = false;
+    Index* storage = nullptr;
 
     /// the index is built or not
-    bool is_built;
+    bool is_built = false;
 
     /// K of KNN graph for building
-    int GK;
+    int GK = 64;
 
     /// indicate how to build a knn graph
     /// - 0: build NSG with brute force search
     /// - 1: build NSG with NNDescent
-    char build_type;
+    char build_type = 0;
 
     /// parameters for nndescent
-    int nndescent_S;
-    int nndescent_R;
-    int nndescent_L;
-    int nndescent_iter;
+    int nndescent_S = 10;
+    int nndescent_R = 100;
+    int nndescent_L; // set to GK + 50
+    int nndescent_iter = 10;
 
     explicit IndexNSG(int d = 0, int R = 32, MetricType metric = METRIC_L2);
     explicit IndexNSG(Index* storage, int R = 32);
@@ -90,7 +90,7 @@ struct IndexNSGFlat : IndexNSG {
  */
 struct IndexNSGPQ : IndexNSG {
     IndexNSGPQ();
-    IndexNSGPQ(int d, int pq_m, int M);
+    IndexNSGPQ(int d, int pq_m, int M, int pq_nbits = 8);
     void train(idx_t n, const float* x) override;
 };
 

data/vendor/faiss/faiss/IndexNeuralNetCodec.cpp

@@ -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.
+ */
+
+#include <faiss/IndexNeuralNetCodec.h>
+#include <faiss/impl/FaissAssert.h>
+#include <faiss/utils/hamming.h>
+
+namespace faiss {
+
+/*********************************************************
+ * IndexNeuralNetCodec implementation
+ *********************************************************/
+
+IndexNeuralNetCodec::IndexNeuralNetCodec(
+        int d,
+        int M,
+        int nbits,
+        MetricType metric)
+        : IndexFlatCodes((M * nbits + 7) / 8, d, metric), M(M), nbits(nbits) {
+    is_trained = false;
+}
+
+void IndexNeuralNetCodec::train(idx_t n, const float* x) {
+    FAISS_THROW_MSG("Training not implemented in C++, use Pytorch");
+}
+
+void IndexNeuralNetCodec::sa_encode(idx_t n, const float* x, uint8_t* codes)
+        const {
+    nn::Tensor2D x_tensor(n, d, x);
+    nn::Int32Tensor2D codes_tensor = net->encode(x_tensor);
+    pack_bitstrings(n, M, nbits, codes_tensor.data(), codes, code_size);
+}
+
+void IndexNeuralNetCodec::sa_decode(idx_t n, const uint8_t* codes, float* x)
+        const {
+    nn::Int32Tensor2D codes_tensor(n, M);
+    unpack_bitstrings(n, M, nbits, codes, code_size, codes_tensor.data());
+    nn::Tensor2D x_tensor = net->decode(codes_tensor);
+    memcpy(x, x_tensor.data(), d * n * sizeof(float));
+}
+
+/*********************************************************
+ * IndexQINeuralNetCodec implementation
+ *********************************************************/
+
+IndexQINCo::IndexQINCo(int d, int M, int nbits, int L, int h, MetricType metric)
+        : IndexNeuralNetCodec(d, M, nbits, metric),
+          qinco(d, 1 << nbits, L, M, h) {
+    net = &qinco;
+}
+
+} // namespace faiss

data/vendor/faiss/faiss/IndexNeuralNetCodec.h

@@ -0,0 +1,49 @@
+/**
+ * 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 <vector>
+
+#include <faiss/IndexFlatCodes.h>
+#include <faiss/utils/NeuralNet.h>
+
+namespace faiss {
+
+struct IndexNeuralNetCodec : IndexFlatCodes {
+    NeuralNetCodec* net = nullptr;
+    size_t M, nbits;
+
+    explicit IndexNeuralNetCodec(
+            int d = 0,
+            int M = 0,
+            int nbits = 0,
+            MetricType metric = METRIC_L2);
+
+    void train(idx_t n, const float* x) override;
+
+    void sa_encode(idx_t n, const float* x, uint8_t* codes) const override;
+    void sa_decode(idx_t n, const uint8_t* codes, float* x) const override;
+
+    ~IndexNeuralNetCodec() {}
+};
+
+struct IndexQINCo : IndexNeuralNetCodec {
+    QINCo qinco;
+
+    IndexQINCo(
+            int d,
+            int M,
+            int nbits,
+            int L,
+            int h,
+            MetricType metric = METRIC_L2);
+
+    ~IndexQINCo() {}
+};
+
+} // namespace faiss

data/vendor/faiss/faiss/IndexPQ.cpp

@@ -14,6 +14,7 @@
 #include <cstring>
 
 #include <algorithm>
+#include <memory>
 
 #include <faiss/impl/DistanceComputer.h>
 #include <faiss/impl/FaissAssert.h>
@@ -86,7 +87,7 @@ struct PQDistanceComputer : FlatCodesDistanceComputer {
         ndis++;
 
         float dis = distance_single_code<PQDecoder>(
-                pq, precomputed_table.data(), code);
+                pq.M, pq.nbits, precomputed_table.data(), code);
         return dis;
     }
 
@@ -198,17 +199,16 @@ void IndexPQ::search(
 
     } else { // code-to-code distances
 
-        uint8_t* q_codes = new uint8_t[n * pq.code_size];
-        ScopeDeleter<uint8_t> del(q_codes);
+        std::unique_ptr<uint8_t[]> q_codes(new uint8_t[n * pq.code_size]);
 
         if (!encode_signs) {
-            pq.compute_codes(x, q_codes, n);
+            pq.compute_codes(x, q_codes.get(), n);
         } else {
             FAISS_THROW_IF_NOT(d == pq.nbits * pq.M);
-            memset(q_codes, 0, n * pq.code_size);
+            memset(q_codes.get(), 0, n * pq.code_size);
             for (size_t i = 0; i < n; i++) {
                 const float* xi = x + i * d;
-                uint8_t* code = q_codes + i * pq.code_size;
+                uint8_t* code = q_codes.get() + i * pq.code_size;
                 for (int j = 0; j < d; j++)
                     if (xi[j] > 0)
                         code[j >> 3] |= 1 << (j & 7);
@@ -219,19 +219,18 @@ void IndexPQ::search(
             float_maxheap_array_t res = {
                     size_t(n), size_t(k), labels, distances};
 
-            pq.search_sdc(q_codes, n, codes.data(), ntotal, &res, true);
+            pq.search_sdc(q_codes.get(), n, codes.data(), ntotal, &res, true);
 
         } else {
-            int* idistances = new int[n * k];
-            ScopeDeleter<int> del(idistances);
+            std::unique_ptr<int[]> idistances(new int[n * k]);
 
             int_maxheap_array_t res = {
-                    size_t(n), size_t(k), labels, idistances};
+                    size_t(n), size_t(k), labels, idistances.get()};
 
             if (search_type == ST_HE) {
                 hammings_knn_hc(
                         &res,
-                        q_codes,
+                        q_codes.get(),
                         codes.data(),
                         ntotal,
                         pq.code_size,
@@ -240,7 +239,7 @@ void IndexPQ::search(
             } else if (search_type == ST_generalized_HE) {
                 generalized_hammings_knn_hc(
                         &res,
-                        q_codes,
+                        q_codes.get(),
                         codes.data(),
                         ntotal,
                         pq.code_size,
@@ -263,21 +262,23 @@ void IndexPQStats::reset() {
 
 IndexPQStats indexPQ_stats;
 
+namespace {
+
 template <class HammingComputer>
-static size_t polysemous_inner_loop(
-        const IndexPQ& index,
+size_t polysemous_inner_loop(
+        const IndexPQ* index,
         const float* dis_table_qi,
         const uint8_t* q_code,
         size_t k,
         float* heap_dis,
         int64_t* heap_ids,
         int ht) {
-    int M = index.pq.M;
-    int code_size = index.pq.code_size;
-    int ksub = index.pq.ksub;
-    size_t ntotal = index.ntotal;
+    int M = index->pq.M;
+    int code_size = index->pq.code_size;
+    int ksub = index->pq.ksub;
+    size_t ntotal = index->ntotal;
 
-    const uint8_t* b_code = index.codes.data();
+    const uint8_t* b_code = index->codes.data();
 
     size_t n_pass_i = 0;
 
@@ -305,6 +306,16 @@ static size_t polysemous_inner_loop(
     return n_pass_i;
 }
 
+struct Run_polysemous_inner_loop {
+    using T = size_t;
+    template <class HammingComputer, class... Types>
+    size_t f(Types... args) {
+        return polysemous_inner_loop<HammingComputer>(args...);
+    }
+};
+
+} // anonymous namespace
+
 void IndexPQ::search_core_polysemous(
         idx_t n,
         const float* x,
@@ -321,22 +332,20 @@ void IndexPQ::search_core_polysemous(
     }
 
     // PQ distance tables
-    float* dis_tables = new float[n * pq.ksub * pq.M];
-    ScopeDeleter<float> del(dis_tables);
-    pq.compute_distance_tables(n, x, dis_tables);
+    std::unique_ptr<float[]> dis_tables(new float[n * pq.ksub * pq.M]);
+    pq.compute_distance_tables(n, x, dis_tables.get());
 
     // Hamming embedding queries
-    uint8_t* q_codes = new uint8_t[n * pq.code_size];
-    ScopeDeleter<uint8_t> del2(q_codes);
+    std::unique_ptr<uint8_t[]> q_codes(new uint8_t[n * pq.code_size]);
 
     if (false) {
-        pq.compute_codes(x, q_codes, n);
+        pq.compute_codes(x, q_codes.get(), n);
     } else {
 #pragma omp parallel for
         for (idx_t qi = 0; qi < n; qi++) {
             pq.compute_code_from_distance_table(
-                    dis_tables + qi * pq.M * pq.ksub,
-                    q_codes + qi * pq.code_size);
+                    dis_tables.get() + qi * pq.M * pq.ksub,
+                    q_codes.get() + qi * pq.code_size);
         }
     }
 
@@ -346,54 +355,33 @@ void IndexPQ::search_core_polysemous(
 
 #pragma omp parallel for reduction(+ : n_pass, bad_code_size)
     for (idx_t qi = 0; qi < n; qi++) {
-        const uint8_t* q_code = q_codes + qi * pq.code_size;
+        const uint8_t* q_code = q_codes.get() + qi * pq.code_size;
 
-        const float* dis_table_qi = dis_tables + qi * pq.M * pq.ksub;
+        const float* dis_table_qi = dis_tables.get() + qi * pq.M * pq.ksub;
 
         int64_t* heap_ids = labels + qi * k;
         float* heap_dis = distances + qi * k;
         maxheap_heapify(k, heap_dis, heap_ids);
 
         if (!generalized_hamming) {
-            switch (pq.code_size) {
-#define DISPATCH(cs)                                          \
-    case cs:                                                  \
-        n_pass += polysemous_inner_loop<HammingComputer##cs>( \
-                *this,                                        \
-                dis_table_qi,                                 \
-                q_code,                                       \
-                k,                                            \
-                heap_dis,                                     \
-                heap_ids,                                     \
-                polysemous_ht);                               \
-        break;
-                DISPATCH(4)
-                DISPATCH(8)
-                DISPATCH(16)
-                DISPATCH(32)
-                DISPATCH(20)
-                default:
-                    if (pq.code_size % 4 == 0) {
-                        n_pass += polysemous_inner_loop<HammingComputerDefault>(
-                                *this,
-                                dis_table_qi,
-                                q_code,
-                                k,
-                                heap_dis,
-                                heap_ids,
-                                polysemous_ht);
-                    } else {
-                        bad_code_size++;
-                    }
-                    break;
-            }
-#undef DISPATCH
+            Run_polysemous_inner_loop r;
+            n_pass += dispatch_HammingComputer(
+                    pq.code_size,
+                    r,
+                    this,
+                    dis_table_qi,
+                    q_code,
+                    k,
+                    heap_dis,
+                    heap_ids,
+                    polysemous_ht);
+
         } else { // generalized hamming
             switch (pq.code_size) {
 #define DISPATCH(cs)                                             \
     case cs:                                                     \
         n_pass += polysemous_inner_loop<GenHammingComputer##cs>( \
-                *this,                                           \
+                this,                                            \
                 dis_table_qi,                                    \
                 q_code,                                          \
                 k,                                               \
@@ -407,7 +395,7 @@ void IndexPQ::search_core_polysemous(
                 default:
                     if (pq.code_size % 8 == 0) {
                         n_pass += polysemous_inner_loop<GenHammingComputerM8>(
-                                *this,
+                                this,
                                 dis_table_qi,
                                 q_code,
                                 k,
@@ -450,12 +438,11 @@ void IndexPQ::sa_decode(idx_t n, const uint8_t* bytes, float* x) const {
 
 void IndexPQ::hamming_distance_table(idx_t n, const float* x, int32_t* dis)
         const {
-    uint8_t* q_codes = new uint8_t[n * pq.code_size];
-    ScopeDeleter<uint8_t> del(q_codes);
+    std::unique_ptr<uint8_t[]> q_codes(new uint8_t[n * pq.code_size]);
 
-    pq.compute_codes(x, q_codes, n);
+    pq.compute_codes(x, q_codes.get(), n);
 
-    hammings(q_codes, codes.data(), n, ntotal, pq.code_size, dis);
+    hammings(q_codes.get(), codes.data(), n, ntotal, pq.code_size, dis);
 }
 
 void IndexPQ::hamming_distance_histogram(
@@ -469,16 +456,15 @@ void IndexPQ::hamming_distance_histogram(
     FAISS_THROW_IF_NOT(pq.nbits == 8);
 
     // Hamming embedding queries
-    uint8_t* q_codes = new uint8_t[n * pq.code_size];
-    ScopeDeleter<uint8_t> del(q_codes);
-    pq.compute_codes(x, q_codes, n);
+    std::unique_ptr<uint8_t[]> q_codes(new uint8_t[n * pq.code_size]);
+    pq.compute_codes(x, q_codes.get(), n);
 
     uint8_t* b_codes;
-    ScopeDeleter<uint8_t> del_b_codes;
+    std::unique_ptr<uint8_t[]> del_b_codes;
 
     if (xb) {
         b_codes = new uint8_t[nb * pq.code_size];
-        del_b_codes.set(b_codes);
+        del_b_codes.reset(b_codes);
         pq.compute_codes(xb, b_codes, nb);
     } else {
         nb = ntotal;
@@ -491,8 +477,7 @@ void IndexPQ::hamming_distance_histogram(
 #pragma omp parallel
     {
         std::vector<int64_t> histi(nbits + 1);
-        hamdis_t* distances = new hamdis_t[nb * bs];
-        ScopeDeleter<hamdis_t> del(distances);
+        std::unique_ptr<hamdis_t[]> distances(new hamdis_t[nb * bs]);
 #pragma omp for
         for (idx_t q0 = 0; q0 < n; q0 += bs) {
             // printf ("dis stats: %zd/%zd\n", q0, n);
@@ -501,12 +486,12 @@ void IndexPQ::hamming_distance_histogram(
                 q1 = n;
 
             hammings(
-                    q_codes + q0 * pq.code_size,
+                    q_codes.get() + q0 * pq.code_size,
                     b_codes,
                     q1 - q0,
                     nb,
                     pq.code_size,
-                    distances);
+                    distances.get());
 
             for (size_t i = 0; i < nb * (q1 - q0); i++)
                 histi[distances[i]]++;
@@ -639,7 +624,7 @@ struct SemiSortedArray {
     int N;
 
     // type of the heap: CMax = sort ascending
-    typedef CMax<T, int> HC;
+    using HC = CMax<T, int>;
     std::vector<int> perm;
 
     int k; // k elements are sorted
@@ -733,7 +718,7 @@ struct MinSumK {
      * We use a heap to maintain a queue of sums, with the associated
      * terms involved in the sum.
      */
-    typedef CMin<T, int64_t> HC;
+    using HC = CMin<T, int64_t>;
     size_t heap_capacity, heap_size;
     T* bh_val;
     int64_t* bh_ids;
@@ -827,7 +812,7 @@ struct MinSumK {
             // enqueue followers
             int64_t ii = ti;
             for (int m = 0; m < M; m++) {
-                int64_t n = ii & ((1L << nbit) - 1);
+                int64_t n = ii & (((int64_t)1 << nbit) - 1);
                 ii >>= nbit;
                 if (n + 1 >= N)
                     continue;
@@ -851,7 +836,7 @@ struct MinSumK {
             }
             int64_t ti = 0;
             for (int m = 0; m < M; m++) {
-                int64_t n = ii & ((1L << nbit) - 1);
+                int64_t n = ii & (((int64_t)1 << nbit) - 1);
                 ti += int64_t(ssx[m].get_ord(n)) << (nbit * m);
                 ii >>= nbit;
             }
@@ -923,17 +908,16 @@ void MultiIndexQuantizer::search(
         return;
     }
 
-    float* dis_tables = new float[n * pq.ksub * pq.M];
-    ScopeDeleter<float> del(dis_tables);
+    std::unique_ptr<float[]> dis_tables(new float[n * pq.ksub * pq.M]);
 
-    pq.compute_distance_tables(n, x, dis_tables);
+    pq.compute_distance_tables(n, x, dis_tables.get());
 
     if (k == 1) {
         // simple version that just finds the min in each table
 
 #pragma omp parallel for
         for (int i = 0; i < n; i++) {
-            const float* dis_table = dis_tables + i * pq.ksub * pq.M;
+            const float* dis_table = dis_tables.get() + i * pq.ksub * pq.M;
             float dis = 0;
             idx_t label = 0;
 
@@ -963,7 +947,7 @@ void MultiIndexQuantizer::search(
                     k, pq.M, pq.nbits, pq.ksub);
 #pragma omp for
             for (int i = 0; i < n; i++) {
-                msk.run(dis_tables + i * pq.ksub * pq.M,
+                msk.run(dis_tables.get() + i * pq.ksub * pq.M,
                         pq.ksub,
                         distances + i * k,
                         labels + i * k);
@@ -975,7 +959,7 @@ void MultiIndexQuantizer::search(
 void MultiIndexQuantizer::reconstruct(idx_t key, float* recons) const {
     int64_t jj = key;
     for (int m = 0; m < pq.M; m++) {
-        int64_t n = jj & ((1L << pq.nbits) - 1);
+        int64_t n = jj & (((int64_t)1 << pq.nbits) - 1);
         jj >>= pq.nbits;
         memcpy(recons, pq.get_centroids(m, n), sizeof(recons[0]) * pq.dsub);
         recons += pq.dsub;
@@ -1107,7 +1091,7 @@ void MultiIndexQuantizer2::search(
 
                 const idx_t* idmap0 = sub_ids.data() + i * k2;
                 int64_t ld_idmap = k2 * n;
-                int64_t mask1 = ksub - 1L;
+                int64_t mask1 = ksub - (int64_t)1;
 
                 for (int k = 0; k < K; k++) {
                     const idx_t* idmap = idmap0;

data/vendor/faiss/faiss/IndexPQ.h

@@ -31,10 +31,7 @@ struct IndexPQ : IndexFlatCodes {
      * @param M      number of subquantizers
      * @param nbits  number of bit per subvector index
      */
-    IndexPQ(int d,        ///< dimensionality of the input vectors
-            size_t M,     ///< number of subquantizers
-            size_t nbits, ///< number of bit per subvector index
-            MetricType metric = METRIC_L2);
+    IndexPQ(int d, size_t M, size_t nbits, MetricType metric = METRIC_L2);
 
     IndexPQ();
 

data/vendor/faiss/faiss/IndexPQFastScan.cpp

@@ -7,8 +7,8 @@
 
 #include <faiss/IndexPQFastScan.h>
 
-#include <limits.h>
 #include <cassert>
+#include <climits>
 #include <memory>
 
 #include <omp.h>