faiss 0.2.4 → 0.2.5

data/vendor/faiss/faiss/impl/ProductQuantizer.cpp

@@ -43,137 +43,12 @@ int sgemm_(
 
 namespace faiss {
 
-/* compute an estimator using look-up tables for typical values of M */
-template <typename CT, class C>
-void pq_estimators_from_tables_Mmul4(
-        int M,
-        const CT* codes,
-        size_t ncodes,
-        const float* __restrict dis_table,
-        size_t ksub,
-        size_t k,
-        float* heap_dis,
-        int64_t* heap_ids) {
-    for (size_t j = 0; j < ncodes; j++) {
-        float dis = 0;
-        const float* dt = dis_table;
-
-        for (size_t m = 0; m < M; m += 4) {
-            float dism = 0;
-            dism = dt[*codes++];
-            dt += ksub;
-            dism += dt[*codes++];
-            dt += ksub;
-            dism += dt[*codes++];
-            dt += ksub;
-            dism += dt[*codes++];
-            dt += ksub;
-            dis += dism;
-        }
-
-        if (C::cmp(heap_dis[0], dis)) {
-            heap_replace_top<C>(k, heap_dis, heap_ids, dis, j);
-        }
-    }
-}
-
-template <typename CT, class C>
-void pq_estimators_from_tables_M4(
-        const CT* codes,
-        size_t ncodes,
-        const float* __restrict dis_table,
-        size_t ksub,
-        size_t k,
-        float* heap_dis,
-        int64_t* heap_ids) {
-    for (size_t j = 0; j < ncodes; j++) {
-        float dis = 0;
-        const float* dt = dis_table;
-        dis = dt[*codes++];
-        dt += ksub;
-        dis += dt[*codes++];
-        dt += ksub;
-        dis += dt[*codes++];
-        dt += ksub;
-        dis += dt[*codes++];
-
-        if (C::cmp(heap_dis[0], dis)) {
-            heap_replace_top<C>(k, heap_dis, heap_ids, dis, j);
-        }
-    }
-}
-
-template <typename CT, class C>
-static inline void pq_estimators_from_tables(
-        const ProductQuantizer& pq,
-        const CT* codes,
-        size_t ncodes,
-        const float* dis_table,
-        size_t k,
-        float* heap_dis,
-        int64_t* heap_ids) {
-    if (pq.M == 4) {
-        pq_estimators_from_tables_M4<CT, C>(
-                codes, ncodes, dis_table, pq.ksub, k, heap_dis, heap_ids);
-        return;
-    }
-
-    if (pq.M % 4 == 0) {
-        pq_estimators_from_tables_Mmul4<CT, C>(
-                pq.M, codes, ncodes, dis_table, pq.ksub, k, heap_dis, heap_ids);
-        return;
-    }
-
-    /* Default is relatively slow */
-    const size_t M = pq.M;
-    const size_t ksub = pq.ksub;
-    for (size_t j = 0; j < ncodes; j++) {
-        float dis = 0;
-        const float* __restrict dt = dis_table;
-        for (int m = 0; m < M; m++) {
-            dis += dt[*codes++];
-            dt += ksub;
-        }
-        if (C::cmp(heap_dis[0], dis)) {
-            heap_replace_top<C>(k, heap_dis, heap_ids, dis, j);
-        }
-    }
-}
-
-template <class C>
-static inline void pq_estimators_from_tables_generic(
-        const ProductQuantizer& pq,
-        size_t nbits,
-        const uint8_t* codes,
-        size_t ncodes,
-        const float* dis_table,
-        size_t k,
-        float* heap_dis,
-        int64_t* heap_ids) {
-    const size_t M = pq.M;
-    const size_t ksub = pq.ksub;
-    for (size_t j = 0; j < ncodes; ++j) {
-        PQDecoderGeneric decoder(codes + j * pq.code_size, nbits);
-        float dis = 0;
-        const float* __restrict dt = dis_table;
-        for (size_t m = 0; m < M; m++) {
-            uint64_t c = decoder.decode();
-            dis += dt[c];
-            dt += ksub;
-        }
-
-        if (C::cmp(heap_dis[0], dis)) {
-            heap_replace_top<C>(k, heap_dis, heap_ids, dis, j);
-        }
-    }
-}
-
 /*********************************************
  * PQ implementation
  *********************************************/
 
 ProductQuantizer::ProductQuantizer(size_t d, size_t M, size_t nbits)
-        : d(d), M(M), nbits(nbits), assign_index(nullptr) {
+        : Quantizer(d, 0), M(M), nbits(nbits), assign_index(nullptr) {
     set_derived_values();
 }
 
@@ -246,7 +121,7 @@ static void init_hypercube_pca(
     }
 }
 
-void ProductQuantizer::train(int n, const float* x) {
+void ProductQuantizer::train(size_t n, const float* x) {
     if (train_type != Train_shared) {
         train_type_t final_train_type;
         final_train_type = train_type;
@@ -321,26 +196,66 @@ void ProductQuantizer::train(int n, const float* x) {
 template <class PQEncoder>
 void compute_code(const ProductQuantizer& pq, const float* x, uint8_t* code) {
     std::vector<float> distances(pq.ksub);
+
+    // It seems to be meaningless to allocate std::vector<float> distances.
+    // But it is done in order to cope the ineffectiveness of the way
+    // the compiler generates the code. Basically, doing something like
+    //
+    //     size_t min_distance = HUGE_VALF;
+    //     size_t idxm = 0;
+    //     for (size_t i = 0; i < N; i++) {
+    //         const float distance = compute_distance(x, y + i * d, d);
+    //         if (distance < min_distance) {
+    //            min_distance = distance;
+    //            idxm = i;
+    //         }
+    //     }
+    //
+    // generates significantly more CPU instructions than the baseline
+    //
+    //     std::vector<float> distances_cached(N);
+    //     for (size_t i = 0; i < N; i++) {
+    //         distances_cached[i] = compute_distance(x, y + i * d, d);
+    //     }
+    //     size_t min_distance = HUGE_VALF;
+    //     size_t idxm = 0;
+    //     for (size_t i = 0; i < N; i++) {
+    //         const float distance = distances_cached[i];
+    //         if (distance < min_distance) {
+    //            min_distance = distance;
+    //            idxm = i;
+    //         }
+    //     }
+    //
+    // So, the baseline is faster. This is because of the vectorization.
+    // I suppose that the branch predictor might affect the performance as well.
+    // So, the buffer is allocated, but it might be unused in
+    // manually optimized code. Let's hope that the compiler is smart enough to
+    // get rid of std::vector allocation in such a case.
+
     PQEncoder encoder(code, pq.nbits);
     for (size_t m = 0; m < pq.M; m++) {
-        float mindis = 1e20;
-        uint64_t idxm = 0;
         const float* xsub = x + m * pq.dsub;
 
-        fvec_L2sqr_ny(
-                distances.data(),
-                xsub,
-                pq.get_centroids(m, 0),
-                pq.dsub,
-                pq.ksub);
-
-        /* Find best centroid */
-        for (size_t i = 0; i < pq.ksub; i++) {
-            float dis = distances[i];
-            if (dis < mindis) {
-                mindis = dis;
-                idxm = i;
-            }
+        uint64_t idxm = 0;
+        if (pq.transposed_centroids.empty()) {
+            // the regular version
+            idxm = fvec_L2sqr_ny_nearest(
+                    distances.data(),
+                    xsub,
+                    pq.get_centroids(m, 0),
+                    pq.dsub,
+                    pq.ksub);
+        } else {
+            // transposed centroids are available, use'em
+            idxm = fvec_L2sqr_ny_nearest_y_transposed(
+                    distances.data(),
+                    xsub,
+                    pq.transposed_centroids.data() + m * pq.ksub,
+                    pq.centroids_sq_lengths.data() + m * pq.ksub,
+                    pq.dsub,
+                    pq.M * pq.ksub,
+                    pq.ksub);
         }
 
         encoder.encode(idxm);
@@ -469,10 +384,13 @@ void ProductQuantizer::compute_codes_with_assign_index(
     }
 }
 
+// block size used in ProductQuantizer::compute_codes
+int product_quantizer_compute_codes_bs = 256 * 1024;
+
 void ProductQuantizer::compute_codes(const float* x, uint8_t* codes, size_t n)
         const {
     // process by blocks to avoid using too much RAM
-    size_t bs = 256 * 1024;
+    size_t bs = product_quantizer_compute_codes_bs;
     if (n > bs) {
         for (size_t i0 = 0; i0 < n; i0 += bs) {
             size_t i1 = std::min(i0 + bs, n);
@@ -606,8 +524,140 @@ void ProductQuantizer::compute_inner_prod_tables(
     }
 }
 
+/**********************************************
+ * Templatized search functions
+ * The template class C indicates whether to keep the highest or smallest values
+ **********************************************/
+
+namespace {
+
+/* compute an estimator using look-up tables for typical values of M */
+template <typename CT, class C>
+void pq_estimators_from_tables_Mmul4(
+        int M,
+        const CT* codes,
+        size_t ncodes,
+        const float* __restrict dis_table,
+        size_t ksub,
+        size_t k,
+        float* heap_dis,
+        int64_t* heap_ids) {
+    for (size_t j = 0; j < ncodes; j++) {
+        float dis = 0;
+        const float* dt = dis_table;
+
+        for (size_t m = 0; m < M; m += 4) {
+            float dism = 0;
+            dism = dt[*codes++];
+            dt += ksub;
+            dism += dt[*codes++];
+            dt += ksub;
+            dism += dt[*codes++];
+            dt += ksub;
+            dism += dt[*codes++];
+            dt += ksub;
+            dis += dism;
+        }
+
+        if (C::cmp(heap_dis[0], dis)) {
+            heap_replace_top<C>(k, heap_dis, heap_ids, dis, j);
+        }
+    }
+}
+
+template <typename CT, class C>
+void pq_estimators_from_tables_M4(
+        const CT* codes,
+        size_t ncodes,
+        const float* __restrict dis_table,
+        size_t ksub,
+        size_t k,
+        float* heap_dis,
+        int64_t* heap_ids) {
+    for (size_t j = 0; j < ncodes; j++) {
+        float dis = 0;
+        const float* dt = dis_table;
+        dis = dt[*codes++];
+        dt += ksub;
+        dis += dt[*codes++];
+        dt += ksub;
+        dis += dt[*codes++];
+        dt += ksub;
+        dis += dt[*codes++];
+
+        if (C::cmp(heap_dis[0], dis)) {
+            heap_replace_top<C>(k, heap_dis, heap_ids, dis, j);
+        }
+    }
+}
+
+template <typename CT, class C>
+void pq_estimators_from_tables(
+        const ProductQuantizer& pq,
+        const CT* codes,
+        size_t ncodes,
+        const float* dis_table,
+        size_t k,
+        float* heap_dis,
+        int64_t* heap_ids) {
+    if (pq.M == 4) {
+        pq_estimators_from_tables_M4<CT, C>(
+                codes, ncodes, dis_table, pq.ksub, k, heap_dis, heap_ids);
+        return;
+    }
+
+    if (pq.M % 4 == 0) {
+        pq_estimators_from_tables_Mmul4<CT, C>(
+                pq.M, codes, ncodes, dis_table, pq.ksub, k, heap_dis, heap_ids);
+        return;
+    }
+
+    /* Default is relatively slow */
+    const size_t M = pq.M;
+    const size_t ksub = pq.ksub;
+    for (size_t j = 0; j < ncodes; j++) {
+        float dis = 0;
+        const float* __restrict dt = dis_table;
+        for (int m = 0; m < M; m++) {
+            dis += dt[*codes++];
+            dt += ksub;
+        }
+        if (C::cmp(heap_dis[0], dis)) {
+            heap_replace_top<C>(k, heap_dis, heap_ids, dis, j);
+        }
+    }
+}
+
+template <class C>
+void pq_estimators_from_tables_generic(
+        const ProductQuantizer& pq,
+        size_t nbits,
+        const uint8_t* codes,
+        size_t ncodes,
+        const float* dis_table,
+        size_t k,
+        float* heap_dis,
+        int64_t* heap_ids) {
+    const size_t M = pq.M;
+    const size_t ksub = pq.ksub;
+    for (size_t j = 0; j < ncodes; ++j) {
+        PQDecoderGeneric decoder(codes + j * pq.code_size, nbits);
+        float dis = 0;
+        const float* __restrict dt = dis_table;
+        for (size_t m = 0; m < M; m++) {
+            uint64_t c = decoder.decode();
+            dis += dt[c];
+            dt += ksub;
+        }
+
+        if (C::cmp(heap_dis[0], dis)) {
+            heap_replace_top<C>(k, heap_dis, heap_ids, dis, j);
+        }
+    }
+}
+
 template <class C>
-static void pq_knn_search_with_tables(
+void pq_knn_search_with_tables(
         const ProductQuantizer& pq,
         size_t nbits,
         const float* dis_tables,
@@ -667,6 +717,8 @@ static void pq_knn_search_with_tables(
     }
 }
 
+} // anonymous namespace
+
 void ProductQuantizer::search(
         const float* __restrict x,
         size_t nx,
@@ -781,4 +833,32 @@ void ProductQuantizer::search_sdc(
     }
 }
 
+void ProductQuantizer::sync_transposed_centroids() {
+    transposed_centroids.resize(d * ksub);
+    centroids_sq_lengths.resize(ksub * M);
+
+    for (size_t mi = 0; mi < M; mi++) {
+        for (size_t ki = 0; ki < ksub; ki++) {
+            float sqlen = 0;
+
+            for (size_t di = 0; di < dsub; di++) {
+                const float q = centroids[(mi * ksub + ki) * dsub + di];
+
+                transposed_centroids[(di * M + mi) * ksub + ki] = q;
+                sqlen += q * q;
+            }
+
+            centroids_sq_lengths[mi * ksub + ki] = sqlen;
+        }
+    }
+}
+
+void ProductQuantizer::clear_transposed_centroids() {
+    transposed_centroids.clear();
+    transposed_centroids.shrink_to_fit();
+
+    centroids_sq_lengths.clear();
+    centroids_sq_lengths.shrink_to_fit();
+}
+
 } // namespace faiss

data/vendor/faiss/faiss/impl/ProductQuantizer.h

@@ -15,23 +15,23 @@
 #include <vector>
 
 #include <faiss/Clustering.h>
+#include <faiss/impl/Quantizer.h>
+#include <faiss/impl/platform_macros.h>
 #include <faiss/utils/Heap.h>
 
 namespace faiss {
 
 /** Product Quantizer. Implemented only for METRIC_L2 */
-struct ProductQuantizer {
+struct ProductQuantizer : Quantizer {
     using idx_t = Index::idx_t;
 
-    size_t d;     ///< size of the input vectors
     size_t M;     ///< number of subquantizers
     size_t nbits; ///< number of bits per quantization index
 
     // values derived from the above
-    size_t dsub;      ///< dimensionality of each subvector
-    size_t code_size; ///< bytes per indexed vector
-    size_t ksub;      ///< number of centroids for each subquantizer
-    bool verbose;     ///< verbose during training?
+    size_t dsub;  ///< dimensionality of each subvector
+    size_t ksub;  ///< number of centroids for each subquantizer
+    bool verbose; ///< verbose during training?
 
     /// initialization
     enum train_type_t {
@@ -49,9 +49,18 @@ struct ProductQuantizer {
     /// d / M)
     Index* assign_index;
 
-    /// Centroid table, size M * ksub * dsub
+    /// Centroid table, size M * ksub * dsub.
+    /// Layout: (M, ksub, dsub)
     std::vector<float> centroids;
 
+    /// Transposed centroid table, size M * ksub * dsub.
+    /// Layout: (dsub, M, ksub)
+    std::vector<float> transposed_centroids;
+
+    /// Squared lengths of centroids, size M * ksub
+    /// Layout: (M, ksub)
+    std::vector<float> centroids_sq_lengths;
+
     /// return the centroids associated with subvector m
     float* get_centroids(size_t m, size_t i) {
         return &centroids[(m * ksub + i) * dsub];
@@ -62,7 +71,7 @@ struct ProductQuantizer {
 
     // Train the product quantizer on a set of points. A clustering
     // can be set on input to define non-default clustering parameters
-    void train(int n, const float* x);
+    void train(size_t n, const float* x) override;
 
     ProductQuantizer(
             size_t d,      /* dimensionality of the input vectors */
@@ -81,7 +90,7 @@ struct ProductQuantizer {
     void compute_code(const float* x, uint8_t* code) const;
 
     /// same as compute_code for several vectors
-    void compute_codes(const float* x, uint8_t* codes, size_t n) const;
+    void compute_codes(const float* x, uint8_t* codes, size_t n) const override;
 
     /// speed up code assignment using assign_index
     /// (non-const because the index is changed)
@@ -92,7 +101,7 @@ struct ProductQuantizer {
 
     /// decode a vector from a given code (or n vectors if third argument)
     void decode(const uint8_t* code, float* x) const;
-    void decode(const uint8_t* code, float* x, size_t n) const;
+    void decode(const uint8_t* code, float* x, size_t n) const override;
 
     /// If we happen to have the distance tables precomputed, this is
     /// more efficient to compute the codes.
@@ -165,8 +174,18 @@ struct ProductQuantizer {
             const size_t ncodes,
             float_maxheap_array_t* res,
             bool init_finalize_heap = true) const;
+
+    /// Sync transposed centroids with regular centroids. This call
+    /// is needed if centroids were edited directly.
+    void sync_transposed_centroids();
+
+    /// Clear transposed centroids table so ones are no longer used.
+    void clear_transposed_centroids();
 };
 
+// block size used in ProductQuantizer::compute_codes
+FAISS_API extern int product_quantizer_compute_codes_bs;
+
 /*************************************************
  * Objects to encode / decode strings of bits
  *************************************************/

data/vendor/faiss/faiss/impl/Quantizer.h

@@ -0,0 +1,43 @@
+// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary.
+
+#pragma once
+
+#include <stdint.h>
+
+namespace faiss {
+
+/** Product Quantizer. Implemented only for METRIC_L2 */
+struct Quantizer {
+    using idx_t = Index::idx_t;
+
+    size_t d;         ///< size of the input vectors
+    size_t code_size; ///< bytes per indexed vector
+
+    explicit Quantizer(size_t d = 0, size_t code_size = 0)
+            : d(d), code_size(code_size) {}
+
+    /** Train the quantizer
+     *
+     * @param x       training vectors, size n * d
+     */
+    virtual void train(size_t n, const float* x) = 0;
+
+    /** Quantize a set of vectors
+     *
+     * @param x        input vectors, size n * d
+     * @param codes    output codes, size n * code_size
+     */
+    virtual void compute_codes(const float* x, uint8_t* codes, size_t n)
+            const = 0;
+
+    /** Decode a set of vectors
+     *
+     * @param codes    input codes, size n * code_size
+     * @param x        output vectors, size n * d
+     */
+    virtual void decode(const uint8_t* code, float* x, size_t n) const = 0;
+
+    virtual ~Quantizer() {}
+};
+
+} // namespace faiss