faiss 0.2.3 → 0.2.4

data/vendor/faiss/faiss/impl/ResidualQuantizer.cpp

@@ -7,18 +7,19 @@
 
 // -*- c++ -*-
 
-#include "faiss/impl/ResidualQuantizer.h"
-#include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/ResidualQuantizer.h>
-#include "faiss/utils/utils.h"
 
+#include <algorithm>
 #include <cstddef>
 #include <cstdio>
 #include <cstring>
 #include <memory>
 
-#include <algorithm>
+#include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/ResidualQuantizer.h>
+#include <faiss/utils/utils.h>
 
+#include <faiss/Clustering.h>
 #include <faiss/IndexFlat.h>
 #include <faiss/VectorTransform.h>
 #include <faiss/impl/AuxIndexStructures.h>
@@ -26,13 +27,34 @@
 #include <faiss/utils/Heap.h>
 #include <faiss/utils/distances.h>
 #include <faiss/utils/hamming.h>
+#include <faiss/utils/simdlib.h>
 #include <faiss/utils/utils.h>
 
+extern "C" {
+
+// general matrix multiplication
+int sgemm_(
+        const char* transa,
+        const char* transb,
+        FINTEGER* m,
+        FINTEGER* n,
+        FINTEGER* k,
+        const float* alpha,
+        const float* a,
+        FINTEGER* lda,
+        const float* b,
+        FINTEGER* ldb,
+        float* beta,
+        float* c,
+        FINTEGER* ldc);
+}
+
 namespace faiss {
 
 ResidualQuantizer::ResidualQuantizer()
         : train_type(Train_progressive_dim),
-          max_beam_size(30),
+          max_beam_size(5),
+          use_beam_LUT(0),
           max_mem_distances(5 * (size_t(1) << 30)), // 5 GiB
           assign_index_factory(nullptr) {
     d = 0;
@@ -40,26 +62,24 @@ ResidualQuantizer::ResidualQuantizer()
     verbose = false;
 }
 
-ResidualQuantizer::ResidualQuantizer(size_t d, const std::vector<size_t>& nbits)
+ResidualQuantizer::ResidualQuantizer(
+        size_t d,
+        const std::vector<size_t>& nbits,
+        Search_type_t search_type)
         : ResidualQuantizer() {
+    this->search_type = search_type;
     this->d = d;
     M = nbits.size();
     this->nbits = nbits;
     set_derived_values();
 }
 
-ResidualQuantizer::ResidualQuantizer(size_t d, size_t M, size_t nbits)
-        : ResidualQuantizer(d, std::vector<size_t>(M, nbits)) {}
-
-namespace {
-
-void fvec_sub(size_t d, const float* a, const float* b, float* c) {
-    for (size_t i = 0; i < d; i++) {
-        c[i] = a[i] - b[i];
-    }
-}
-
-} // anonymous namespace
+ResidualQuantizer::ResidualQuantizer(
+        size_t d,
+        size_t M,
+        size_t nbits,
+        Search_type_t search_type)
+        : ResidualQuantizer(d, std::vector<size_t>(M, nbits), search_type) {}
 
 void beam_search_encode_step(
         size_t d,
@@ -90,7 +110,7 @@ void beam_search_encode_step(
         cent_ids.resize(n * beam_size * new_beam_size);
         if (assign_index->ntotal != 0) {
             // then we assume the codebooks are already added to the index
-            FAISS_THROW_IF_NOT(assign_index->ntotal != K);
+            FAISS_THROW_IF_NOT(assign_index->ntotal == K);
         } else {
             assign_index->add(K, cent);
         }
@@ -208,6 +228,7 @@ void ResidualQuantizer::train(size_t n, const float* x) {
     std::vector<int32_t> codes;
     std::vector<float> distances;
     double t0 = getmillisecs();
+    double clustering_time = 0;
 
     for (int m = 0; m < M; m++) {
         int K = 1 << nbits[m];
@@ -224,7 +245,7 @@ void ResidualQuantizer::train(size_t n, const float* x) {
             }
             train_residuals = residuals1;
         }
-        train_type_t tt = train_type_t(train_type & ~Train_top_beam);
+        train_type_t tt = train_type_t(train_type & 1023);
 
         std::vector<float> codebooks;
         float obj = 0;
@@ -235,6 +256,9 @@ void ResidualQuantizer::train(size_t n, const float* x) {
         } else {
             assign_index.reset(new IndexFlatL2(d));
         }
+
+        double t1 = getmillisecs();
+
         if (tt == Train_default) {
             Clustering clus(d, K, cp);
             clus.train(
@@ -256,6 +280,7 @@ void ResidualQuantizer::train(size_t n, const float* x) {
         } else {
             FAISS_THROW_MSG("train type not supported");
         }
+        clustering_time += (getmillisecs() - t1) / 1000;
 
         memcpy(this->codebooks.data() + codebook_offsets[m] * d,
                codebooks.data(),
@@ -268,21 +293,38 @@ void ResidualQuantizer::train(size_t n, const float* x) {
         std::vector<float> new_residuals(n * new_beam_size * d);
         std::vector<float> new_distances(n * new_beam_size);
 
-        beam_search_encode_step(
-                d,
-                K,
-                codebooks.data(),
-                n,
-                cur_beam_size,
-                residuals.data(),
-                m,
-                codes.data(),
-                new_beam_size,
-                new_codes.data(),
-                new_residuals.data(),
-                new_distances.data(),
-                assign_index.get());
+        size_t bs;
+        { // determine batch size
+            size_t mem = memory_per_point();
+            if (n > 1 && mem * n > max_mem_distances) {
+                // then split queries to reduce temp memory
+                bs = std::max(max_mem_distances / mem, size_t(1));
+            } else {
+                bs = n;
+            }
+        }
 
+        for (size_t i0 = 0; i0 < n; i0 += bs) {
+            size_t i1 = std::min(i0 + bs, n);
+
+            /* printf("i0: %ld i1: %ld K %d ntotal assign index %ld\n",
+                i0, i1, K, assign_index->ntotal); */
+
+            beam_search_encode_step(
+                    d,
+                    K,
+                    codebooks.data(),
+                    i1 - i0,
+                    cur_beam_size,
+                    residuals.data() + i0 * cur_beam_size * d,
+                    m,
+                    codes.data() + i0 * cur_beam_size * m,
+                    new_beam_size,
+                    new_codes.data() + i0 * new_beam_size * (m + 1),
+                    new_residuals.data() + i0 * new_beam_size * d,
+                    new_distances.data() + i0 * new_beam_size,
+                    assign_index.get());
+        }
         codes.swap(new_codes);
         residuals.swap(new_residuals);
         distances.swap(new_distances);
@@ -293,20 +335,57 @@ void ResidualQuantizer::train(size_t n, const float* x) {
         }
 
         if (verbose) {
-            printf("[%.3f s] train stage %d, %d bits, kmeans objective %g, "
-                   "total distance %g, beam_size %d->%d\n",
+            printf("[%.3f s, %.3f s clustering] train stage %d, %d bits, kmeans objective %g, "
+                   "total distance %g, beam_size %d->%d (batch size %zd)\n",
                    (getmillisecs() - t0) / 1000,
+                   clustering_time,
                    m,
                    int(nbits[m]),
                    obj,
                    sum_distances,
                    cur_beam_size,
-                   new_beam_size);
+                   new_beam_size,
+                   bs);
         }
         cur_beam_size = new_beam_size;
     }
 
+    // find min and max norms
+    std::vector<float> norms(n);
+
+    for (size_t i = 0; i < n; i++) {
+        norms[i] = fvec_L2sqr(
+                x + i * d, residuals.data() + i * cur_beam_size * d, d);
+    }
+
+    // fvec_norms_L2sqr(norms.data(), x, d, n);
+
+    norm_min = HUGE_VALF;
+    norm_max = -HUGE_VALF;
+    for (idx_t i = 0; i < n; i++) {
+        if (norms[i] < norm_min) {
+            norm_min = norms[i];
+        }
+        if (norms[i] > norm_max) {
+            norm_max = norms[i];
+        }
+    }
+
+    if (search_type == ST_norm_cqint8 || search_type == ST_norm_cqint4) {
+        size_t k = (1 << 8);
+        if (search_type == ST_norm_cqint4) {
+            k = (1 << 4);
+        }
+        Clustering1D clus(k);
+        clus.train_exact(n, norms.data());
+        qnorm.add(clus.k, clus.centroids.data());
+    }
+
     is_trained = true;
+
+    if (!(train_type & Skip_codebook_tables)) {
+        compute_codebook_tables();
+    }
 }
 
 size_t ResidualQuantizer::memory_per_point(int beam_size) const {
@@ -341,22 +420,76 @@ void ResidualQuantizer::compute_codes(
         return;
     }
 
-    std::vector<float> residuals(max_beam_size * n * d);
     std::vector<int32_t> codes(max_beam_size * M * n);
+    std::vector<float> norms;
     std::vector<float> distances(max_beam_size * n);
 
-    refine_beam(
-            n,
-            1,
-            x,
-            max_beam_size,
-            codes.data(),
-            residuals.data(),
-            distances.data());
+    if (use_beam_LUT == 0) {
+        std::vector<float> residuals(max_beam_size * n * d);
+
+        refine_beam(
+                n,
+                1,
+                x,
+                max_beam_size,
+                codes.data(),
+                residuals.data(),
+                distances.data());
+
+        if (search_type == ST_norm_float || search_type == ST_norm_qint8 ||
+            search_type == ST_norm_qint4) {
+            norms.resize(n);
+            // recover the norms of reconstruction as
+            // || original_vector - residual ||^2
+            for (size_t i = 0; i < n; i++) {
+                norms[i] = fvec_L2sqr(
+                        x + i * d, residuals.data() + i * max_beam_size * d, d);
+            }
+        }
+    } else if (use_beam_LUT == 1) {
+        FAISS_THROW_IF_NOT_MSG(
+                codebook_cross_products.size() ==
+                        total_codebook_size * total_codebook_size,
+                "call compute_codebook_tables first");
+
+        std::vector<float> query_norms(n);
+        fvec_norms_L2sqr(query_norms.data(), x, d, n);
+
+        std::vector<float> query_cp(n * total_codebook_size);
+        {
+            FINTEGER ti = total_codebook_size, di = d, ni = n;
+            float zero = 0, one = 1;
+            sgemm_("Transposed",
+                   "Not transposed",
+                   &ti,
+                   &ni,
+                   &di,
+                   &one,
+                   codebooks.data(),
+                   &di,
+                   x,
+                   &di,
+                   &zero,
+                   query_cp.data(),
+                   &ti);
+        }
 
+        refine_beam_LUT(
+                n,
+                query_norms.data(),
+                query_cp.data(),
+                max_beam_size,
+                codes.data(),
+                distances.data());
+    }
     // pack only the first code of the beam (hence the ld_codes=M *
     // max_beam_size)
-    pack_codes(n, codes.data(), codes_out, M * max_beam_size);
+    pack_codes(
+            n,
+            codes.data(),
+            codes_out,
+            M * max_beam_size,
+            norms.size() > 0 ? norms.data() : nullptr);
 }
 
 void ResidualQuantizer::refine_beam(
@@ -445,4 +578,181 @@ void ResidualQuantizer::refine_beam(
     }
 }
 
+/*******************************************************************
+ * Functions using the dot products between codebook entries
+ *******************************************************************/
+
+void ResidualQuantizer::compute_codebook_tables() {
+    codebook_cross_products.resize(total_codebook_size * total_codebook_size);
+    cent_norms.resize(total_codebook_size);
+    // stricly speaking we could use ssyrk
+    {
+        FINTEGER ni = total_codebook_size;
+        FINTEGER di = d;
+        float zero = 0, one = 1;
+        sgemm_("Transposed",
+               "Not transposed",
+               &ni,
+               &ni,
+               &di,
+               &one,
+               codebooks.data(),
+               &di,
+               codebooks.data(),
+               &di,
+               &zero,
+               codebook_cross_products.data(),
+               &ni);
+    }
+    for (size_t i = 0; i < total_codebook_size; i++) {
+        cent_norms[i] = codebook_cross_products[i + i * total_codebook_size];
+    }
+}
+
+void beam_search_encode_step_tab(
+        size_t K,
+        size_t n,
+        size_t beam_size,                  // input sizes
+        const float* codebook_cross_norms, // size K * ldc
+        size_t ldc,                        // >= K
+        const uint64_t* codebook_offsets,  // m
+        const float* query_cp,             // size n * ldqc
+        size_t ldqc,                       // >= K
+        const float* cent_norms_i,         // size K
+        size_t m,
+        const int32_t* codes,   // n * beam_size * m
+        const float* distances, // n * beam_size
+        size_t new_beam_size,
+        int32_t* new_codes,   // n * new_beam_size * (m + 1)
+        float* new_distances) // n * new_beam_size
+{
+    FAISS_THROW_IF_NOT(ldc >= K);
+
+#pragma omp parallel for if (n > 100)
+    for (int64_t i = 0; i < n; i++) {
+        std::vector<float> cent_distances(beam_size * K);
+        std::vector<float> cd_common(K);
+
+        const int32_t* codes_i = codes + i * m * beam_size;
+        const float* query_cp_i = query_cp + i * ldqc;
+        const float* distances_i = distances + i * beam_size;
+
+        for (size_t k = 0; k < K; k++) {
+            cd_common[k] = cent_norms_i[k] - 2 * query_cp_i[k];
+        }
+
+        for (size_t b = 0; b < beam_size; b++) {
+            std::vector<float> dp(K);
+
+            for (size_t m1 = 0; m1 < m; m1++) {
+                size_t c = codes_i[b * m + m1];
+                const float* cb =
+                        &codebook_cross_norms[(codebook_offsets[m1] + c) * ldc];
+                fvec_add(K, cb, dp.data(), dp.data());
+            }
+
+            for (size_t k = 0; k < K; k++) {
+                cent_distances[b * K + k] =
+                        distances_i[b] + cd_common[k] + 2 * dp[k];
+            }
+        }
+
+        using C = CMax<float, int>;
+        int32_t* new_codes_i = new_codes + i * (m + 1) * new_beam_size;
+        float* new_distances_i = new_distances + i * new_beam_size;
+
+        const float* cent_distances_i = cent_distances.data();
+
+        // then we have to select the best results
+        for (int i = 0; i < new_beam_size; i++) {
+            new_distances_i[i] = C::neutral();
+        }
+        std::vector<int> perm(new_beam_size, -1);
+        heap_addn<C>(
+                new_beam_size,
+                new_distances_i,
+                perm.data(),
+                cent_distances_i,
+                nullptr,
+                beam_size * K);
+        heap_reorder<C>(new_beam_size, new_distances_i, perm.data());
+
+        for (int j = 0; j < new_beam_size; j++) {
+            int js = perm[j] / K;
+            int ls = perm[j] % K;
+            if (m > 0) {
+                memcpy(new_codes_i, codes_i + js * m, sizeof(*codes) * m);
+            }
+            new_codes_i[m] = ls;
+            new_codes_i += m + 1;
+        }
+    }
+}
+
+void ResidualQuantizer::refine_beam_LUT(
+        size_t n,
+        const float* query_norms, // size n
+        const float* query_cp,    //
+        int out_beam_size,
+        int32_t* out_codes,
+        float* out_distances) const {
+    int beam_size = 1;
+
+    std::vector<int32_t> codes;
+    std::vector<float> distances(query_norms, query_norms + n);
+    double t0 = getmillisecs();
+
+    for (int m = 0; m < M; m++) {
+        int K = 1 << nbits[m];
+
+        int new_beam_size = std::min(beam_size * K, out_beam_size);
+        std::vector<int32_t> new_codes(n * new_beam_size * (m + 1));
+        std::vector<float> new_distances(n * new_beam_size);
+
+        beam_search_encode_step_tab(
+                K,
+                n,
+                beam_size,
+                codebook_cross_products.data() + codebook_offsets[m],
+                total_codebook_size,
+                codebook_offsets.data(),
+                query_cp + codebook_offsets[m],
+                total_codebook_size,
+                cent_norms.data() + codebook_offsets[m],
+                m,
+                codes.data(),
+                distances.data(),
+                new_beam_size,
+                new_codes.data(),
+                new_distances.data());
+
+        codes.swap(new_codes);
+        distances.swap(new_distances);
+        beam_size = new_beam_size;
+
+        if (verbose) {
+            float sum_distances = 0;
+            for (int j = 0; j < distances.size(); j++) {
+                sum_distances += distances[j];
+            }
+            printf("[%.3f s] encode stage %d, %d bits, "
+                   "total error %g, beam_size %d\n",
+                   (getmillisecs() - t0) / 1000,
+                   m,
+                   int(nbits[m]),
+                   sum_distances,
+                   beam_size);
+        }
+    }
+
+    if (out_codes) {
+        memcpy(out_codes, codes.data(), codes.size() * sizeof(codes[0]));
+    }
+    if (out_distances) {
+        memcpy(out_distances,
+               distances.data(),
+               distances.size() * sizeof(distances[0]));
+    }
+}
+
 } // namespace faiss

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

@@ -25,18 +25,32 @@ namespace faiss {
 struct ResidualQuantizer : AdditiveQuantizer {
     /// initialization
     enum train_type_t {
-        Train_default,         ///< regular k-means
-        Train_progressive_dim, ///< progressive dim clustering
+        Train_default = 0,         ///< regular k-means
+        Train_progressive_dim = 1, ///< progressive dim clustering
+        Train_default_Train_top_beam = 1024,
+        Train_progressive_dim_Train_top_beam = 1025,
+        Train_default_Skip_codebook_tables = 2048,
+        Train_progressive_dim_Skip_codebook_tables = 2049,
+        Train_default_Train_top_beam_Skip_codebook_tables = 3072,
+        Train_progressive_dim_Train_top_beam_Skip_codebook_tables = 3073,
     };
 
+    train_type_t train_type;
+
     // 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;
+
+    // set this bit to not autmatically compute the codebook tables
+    // after training
+    static const int Skip_codebook_tables = 2048;
 
     /// beam size used for training and for encoding
     int max_beam_size;
 
+    /// use LUT for beam search
+    int use_beam_LUT;
+
     /// distance matrixes with beam search can get large, so use this
     /// to batch computations at encoding time.
     size_t max_mem_distances;
@@ -47,12 +61,16 @@ struct ResidualQuantizer : AdditiveQuantizer {
     /// 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,
+            const std::vector<size_t>& nbits,
+            Search_type_t search_type = ST_decompress);
 
     ResidualQuantizer(
-            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);
 
     ResidualQuantizer();
 
@@ -85,12 +103,32 @@ struct ResidualQuantizer : AdditiveQuantizer {
             float* new_residuals = nullptr,
             float* new_distances = nullptr) const;
 
+    void refine_beam_LUT(
+            size_t n,
+            const float* query_norms,
+            const float* query_cp,
+            int new_beam_size,
+            int32_t* new_codes,
+            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;
+
+    /** Cross products used in codebook tables
+     *
+     * These are used to keep trak of norms of centroids.
+     */
+    void compute_codebook_tables();
+
+    /// dot products of all codebook vectors with each other
+    /// size total_codebook_size * total_codebook_size
+    std::vector<float> codebook_cross_products;
+    /// norms of all vectors
+    std::vector<float> cent_norms;
 };
 
 /** Encode a residual by sampling from a centroid table.
@@ -127,4 +165,24 @@ void beam_search_encode_step(
         float* new_distances,
         Index* assign_index = nullptr);
 
+/** Encode a set of vectors using their dot products with the codebooks
+ *
+ */
+void beam_search_encode_step_tab(
+        size_t K,
+        size_t n,
+        size_t beam_size,                  // input sizes
+        const float* codebook_cross_norms, // size K * ldc
+        size_t ldc,                        // >= K
+        const uint64_t* codebook_offsets,  // m
+        const float* query_cp,             // size n * ldqc
+        size_t ldqc,                       // >= K
+        const float* cent_norms_i,         // size K
+        size_t m,
+        const int32_t* codes,   // n * beam_size * m
+        const float* distances, // n * beam_size
+        size_t new_beam_size,
+        int32_t* new_codes,    // n * new_beam_size * (m + 1)
+        float* new_distances); // n * new_beam_size
+
 }; // namespace faiss

data/vendor/faiss/faiss/impl/ScalarQuantizer.cpp

@@ -1335,12 +1335,9 @@ namespace {
 template <class DCClass>
 struct IVFSQScannerIP : InvertedListScanner {
     DCClass dc;
-    bool store_pairs, by_residual;
+    bool by_residual;
 
-    size_t code_size;
-
-    idx_t list_no; /// current list (set to 0 for Flat index
-    float accu0;   /// added to all distances
+    float accu0; /// added to all distances
 
     IVFSQScannerIP(
             int d,
@@ -1348,12 +1345,10 @@ struct IVFSQScannerIP : InvertedListScanner {
             size_t code_size,
             bool store_pairs,
             bool by_residual)
-            : dc(d, trained),
-              store_pairs(store_pairs),
-              by_residual(by_residual),
-              code_size(code_size),
-              list_no(0),
-              accu0(0) {}
+            : dc(d, trained), by_residual(by_residual), accu0(0) {
+        this->store_pairs = store_pairs;
+        this->code_size = code_size;
+    }
 
     void set_query(const float* query) override {
         dc.set_query(query);
@@ -1411,10 +1406,8 @@ template <class DCClass>
 struct IVFSQScannerL2 : InvertedListScanner {
     DCClass dc;
 
-    bool store_pairs, by_residual;
-    size_t code_size;
+    bool by_residual;
     const Index* quantizer;
-    idx_t list_no;  /// current inverted list
     const float* x; /// current query
 
     std::vector<float> tmp;
@@ -1427,13 +1420,13 @@ struct IVFSQScannerL2 : InvertedListScanner {
             bool store_pairs,
             bool by_residual)
             : dc(d, trained),
-              store_pairs(store_pairs),
               by_residual(by_residual),
-              code_size(code_size),
               quantizer(quantizer),
-              list_no(0),
               x(nullptr),
-              tmp(d) {}
+              tmp(d) {
+        this->store_pairs = store_pairs;
+        this->code_size = code_size;
+    }
 
     void set_query(const float* query) override {
         x = query;
@@ -1443,8 +1436,8 @@ struct IVFSQScannerL2 : InvertedListScanner {
     }
 
     void set_list(idx_t list_no, float /*coarse_dis*/) override {
+        this->list_no = list_no;
         if (by_residual) {
-            this->list_no = list_no;
             // shift of x_in wrt centroid
             quantizer->compute_residual(x, tmp.data(), list_no);
             dc.set_query(tmp.data());