faiss 0.2.3 → 0.2.5

data/vendor/faiss/faiss/gpu/utils/DeviceUtils.h

@@ -70,10 +70,11 @@ bool getTensorCoreSupport(int device);
 /// Equivalent to getTensorCoreSupport(getCurrentDevice())
 bool getTensorCoreSupportCurrentDevice();
 
-/// Returns the maximum k-selection value supported based on the CUDA SDK that
-/// we were compiled with. .cu files can use DeviceDefs.cuh, but this is for
-/// non-CUDA files
-int getMaxKSelection();
+/// Returns the amount of currently available memory on the given device
+size_t getFreeMemory(int device);
+
+/// Equivalent to getFreeMemory(getCurrentDevice())
+size_t getFreeMemoryCurrentDevice();
 
 /// RAII object to set the current device, and restore the previous
 /// device upon destruction

data/vendor/faiss/faiss/impl/AdditiveQuantizer.cpp

@@ -8,7 +8,6 @@
 // -*- c++ -*-
 
 #include <faiss/impl/AdditiveQuantizer.h>
-#include <faiss/impl/FaissAssert.h>
 
 #include <cstddef>
 #include <cstdio>
@@ -18,9 +17,13 @@
 
 #include <algorithm>
 
+#include <faiss/Clustering.h>
+#include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/LocalSearchQuantizer.h>
+#include <faiss/impl/ResidualQuantizer.h>
 #include <faiss/utils/Heap.h>
 #include <faiss/utils/distances.h>
-#include <faiss/utils/hamming.h> // BitstringWriter
+#include <faiss/utils/hamming.h>
 #include <faiss/utils/utils.h>
 
 extern "C" {
@@ -42,51 +45,211 @@ int sgemm_(
         FINTEGER* ldc);
 }
 
-namespace {
-
-// c and a and b can overlap
-void fvec_add(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];
-    }
-}
+namespace faiss {
 
-void fvec_add(size_t d, const float* a, float b, float* c) {
-    for (size_t i = 0; i < d; i++) {
-        c[i] = a[i] + b;
-    }
+AdditiveQuantizer::AdditiveQuantizer(
+        size_t d,
+        const std::vector<size_t>& nbits,
+        Search_type_t search_type)
+        : Quantizer(d),
+          M(nbits.size()),
+          nbits(nbits),
+          verbose(false),
+          is_trained(false),
+          max_mem_distances(5 * (size_t(1) << 30)), // 5 GiB
+          search_type(search_type) {
+    norm_max = norm_min = NAN;
+    tot_bits = 0;
+    total_codebook_size = 0;
+    only_8bit = false;
+    set_derived_values();
 }
 
-} // namespace
-
-namespace faiss {
+AdditiveQuantizer::AdditiveQuantizer()
+        : AdditiveQuantizer(0, std::vector<size_t>()) {}
 
 void AdditiveQuantizer::set_derived_values() {
     tot_bits = 0;
-    is_byte_aligned = true;
+    only_8bit = true;
     codebook_offsets.resize(M + 1, 0);
     for (int i = 0; i < M; i++) {
         int nbit = nbits[i];
         size_t k = 1 << nbit;
         codebook_offsets[i + 1] = codebook_offsets[i] + k;
         tot_bits += nbit;
-        if (nbit % 8 != 0) {
-            is_byte_aligned = false;
+        if (nbit != 0) {
+            only_8bit = false;
         }
     }
     total_codebook_size = codebook_offsets[M];
+    switch (search_type) {
+        case ST_norm_float:
+            norm_bits = 32;
+            break;
+        case ST_norm_qint8:
+        case ST_norm_cqint8:
+        case ST_norm_lsq2x4:
+        case ST_norm_rq2x4:
+            norm_bits = 8;
+            break;
+        case ST_norm_qint4:
+        case ST_norm_cqint4:
+            norm_bits = 4;
+            break;
+        case ST_decompress:
+        case ST_LUT_nonorm:
+        case ST_norm_from_LUT:
+        default:
+            norm_bits = 0;
+            break;
+    }
+    tot_bits += norm_bits;
+
     // convert bits to bytes
     code_size = (tot_bits + 7) / 8;
 }
 
+void AdditiveQuantizer::train_norm(size_t n, const float* norms) {
+    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);
+        qnorm.add(clus.k, clus.centroids.data());
+    } else if (search_type == ST_norm_lsq2x4 || search_type == ST_norm_rq2x4) {
+        std::unique_ptr<AdditiveQuantizer> aq;
+        if (search_type == ST_norm_lsq2x4) {
+            aq.reset(new LocalSearchQuantizer(1, 2, 4));
+        } else {
+            aq.reset(new ResidualQuantizer(1, 2, 4));
+        }
+
+        aq->train(n, norms);
+        // flatten aq codebooks
+        std::vector<float> flat_codebooks(1 << 8);
+        FAISS_THROW_IF_NOT(aq->codebooks.size() == 32);
+
+        // save norm tables for 4-bit fastscan search
+        norm_tabs = aq->codebooks;
+
+        // assume big endian
+        const float* c = norm_tabs.data();
+        for (size_t i = 0; i < 16; i++) {
+            for (size_t j = 0; j < 16; j++) {
+                flat_codebooks[i * 16 + j] = c[j] + c[16 + i];
+            }
+        }
+
+        qnorm.reset();
+        qnorm.add(1 << 8, flat_codebooks.data());
+        FAISS_THROW_IF_NOT(qnorm.ntotal == (1 << 8));
+    }
+}
+
+namespace {
+
+// TODO
+// https://stackoverflow.com/questions/31631224/hacks-for-clamping-integer-to-0-255-and-doubles-to-0-0-1-0
+
+uint8_t encode_qint8(float x, float amin, float amax) {
+    float x1 = (x - amin) / (amax - amin) * 256;
+    int32_t xi = int32_t(floor(x1));
+
+    return xi < 0 ? 0 : xi > 255 ? 255 : xi;
+}
+
+uint8_t encode_qint4(float x, float amin, float amax) {
+    float x1 = (x - amin) / (amax - amin) * 16;
+    int32_t xi = int32_t(floor(x1));
+
+    return xi < 0 ? 0 : xi > 15 ? 15 : xi;
+}
+
+float decode_qint8(uint8_t i, float amin, float amax) {
+    return (i + 0.5) / 256 * (amax - amin) + amin;
+}
+
+float decode_qint4(uint8_t i, float amin, float amax) {
+    return (i + 0.5) / 16 * (amax - amin) + amin;
+}
+
+} // anonymous namespace
+
+uint32_t AdditiveQuantizer::encode_qcint(float x) const {
+    idx_t id;
+    qnorm.assign(1, &x, &id, 1);
+    return uint32_t(id);
+}
+
+float AdditiveQuantizer::decode_qcint(uint32_t c) const {
+    return qnorm.get_xb()[c];
+}
+
+uint64_t AdditiveQuantizer::encode_norm(float norm) const {
+    switch (search_type) {
+        case ST_norm_float:
+            uint32_t inorm;
+            memcpy(&inorm, &norm, 4);
+            return inorm;
+        case ST_norm_qint8:
+            return encode_qint8(norm, norm_min, norm_max);
+        case ST_norm_qint4:
+            return encode_qint4(norm, norm_min, norm_max);
+        case ST_norm_lsq2x4:
+        case ST_norm_rq2x4:
+        case ST_norm_cqint8:
+            return encode_qcint(norm);
+        case ST_norm_cqint4:
+            return encode_qcint(norm);
+        case ST_decompress:
+        case ST_LUT_nonorm:
+        case ST_norm_from_LUT:
+        default:
+            return 0;
+    }
+}
+
 void AdditiveQuantizer::pack_codes(
         size_t n,
         const int32_t* codes,
         uint8_t* packed_codes,
-        int64_t ld_codes) const {
+        int64_t ld_codes,
+        const float* norms,
+        const float* centroids) const {
     if (ld_codes == -1) {
         ld_codes = M;
     }
+    std::vector<float> norm_buf;
+    if (search_type == ST_norm_float || search_type == ST_norm_qint4 ||
+        search_type == ST_norm_qint8 || search_type == ST_norm_cqint8 ||
+        search_type == ST_norm_cqint4 || search_type == ST_norm_lsq2x4 ||
+        search_type == ST_norm_rq2x4) {
+        if (centroids != nullptr || !norms) {
+            norm_buf.resize(n);
+            std::vector<float> x_recons(n * d);
+            decode_unpacked(codes, x_recons.data(), n, ld_codes);
+
+            if (centroids != nullptr) {
+                // x = x + c
+                fvec_add(n * d, x_recons.data(), centroids, x_recons.data());
+            }
+            fvec_norms_L2sqr(norm_buf.data(), x_recons.data(), d, n);
+            norms = norm_buf.data();
+        }
+    }
 #pragma omp parallel for if (n > 1000)
     for (int64_t i = 0; i < n; i++) {
         const int32_t* codes1 = codes + i * ld_codes;
@@ -94,6 +257,9 @@ void AdditiveQuantizer::pack_codes(
         for (int m = 0; m < M; m++) {
             bsw.write(codes1[m], nbits[m]);
         }
+        if (norm_bits != 0) {
+            bsw.write(encode_norm(norms[i]), norm_bits);
+        }
     }
 }
 
@@ -118,10 +284,39 @@ void AdditiveQuantizer::decode(const uint8_t* code, float* x, size_t n) const {
     }
 }
 
+void AdditiveQuantizer::decode_unpacked(
+        const int32_t* code,
+        float* x,
+        size_t n,
+        int64_t ld_codes) const {
+    FAISS_THROW_IF_NOT_MSG(
+            is_trained, "The additive quantizer is not trained yet.");
+
+    if (ld_codes == -1) {
+        ld_codes = M;
+    }
+
+    // standard additive quantizer decoding
+#pragma omp parallel for if (n > 1000)
+    for (int64_t i = 0; i < n; i++) {
+        const int32_t* codesi = code + i * ld_codes;
+        float* xi = x + i * d;
+        for (int m = 0; m < M; m++) {
+            int idx = codesi[m];
+            const float* c = codebooks.data() + d * (codebook_offsets[m] + idx);
+            if (m == 0) {
+                memcpy(xi, c, sizeof(*x) * d);
+            } else {
+                fvec_add(d, xi, c, xi);
+            }
+        }
+    }
+}
+
 AdditiveQuantizer::~AdditiveQuantizer() {}
 
 /****************************************************************************
- * Support for fast distance computations and search with additive quantizer
+ * Support for fast distance computations in centroids
  ****************************************************************************/
 
 void AdditiveQuantizer::compute_centroid_norms(float* norms) const {
@@ -151,28 +346,33 @@ void AdditiveQuantizer::decode_64bit(idx_t bits, float* xi) const {
     }
 }
 
-void AdditiveQuantizer::compute_LUT(size_t n, const float* xq, float* LUT)
-        const {
+void AdditiveQuantizer::compute_LUT(
+        size_t n,
+        const float* xq,
+        float* LUT,
+        float alpha,
+        long ld_lut) const {
     // in all cases, it is large matrix multiplication
 
     FINTEGER ncenti = total_codebook_size;
     FINTEGER di = d;
     FINTEGER nqi = n;
-    float one = 1, zero = 0;
+    FINTEGER ldc = ld_lut > 0 ? ld_lut : ncenti;
+    float zero = 0;
 
     sgemm_("Transposed",
            "Not transposed",
            &ncenti,
            &nqi,
            &di,
-           &one,
+           &alpha,
            codebooks.data(),
            &di,
            xq,
            &di,
            &zero,
            LUT,
-           &ncenti);
+           &ldc);
 }
 
 namespace {
@@ -201,7 +401,7 @@ void compute_inner_prod_with_LUT(
 
 } // anonymous namespace
 
-void AdditiveQuantizer::knn_exact_inner_product(
+void AdditiveQuantizer::knn_centroids_inner_product(
         idx_t n,
         const float* xq,
         idx_t k,
@@ -227,7 +427,7 @@ void AdditiveQuantizer::knn_exact_inner_product(
     }
 }
 
-void AdditiveQuantizer::knn_exact_L2(
+void AdditiveQuantizer::knn_centroids_L2(
         idx_t n,
         const float* xq,
         idx_t k,
@@ -267,4 +467,106 @@ void AdditiveQuantizer::knn_exact_L2(
     }
 }
 
+/****************************************************************************
+ * Support for fast distance computations in codes
+ ****************************************************************************/
+
+namespace {
+
+float accumulate_IPs(
+        const AdditiveQuantizer& aq,
+        BitstringReader& bs,
+        const uint8_t* codes,
+        const float* LUT) {
+    float accu = 0;
+    for (int m = 0; m < aq.M; m++) {
+        size_t nbit = aq.nbits[m];
+        int idx = bs.read(nbit);
+        accu += LUT[idx];
+        LUT += (uint64_t)1 << nbit;
+    }
+    return accu;
+}
+
+} // anonymous namespace
+
+template <>
+float AdditiveQuantizer::
+        compute_1_distance_LUT<true, AdditiveQuantizer::ST_LUT_nonorm>(
+                const uint8_t* codes,
+                const float* LUT) const {
+    BitstringReader bs(codes, code_size);
+    return accumulate_IPs(*this, bs, codes, LUT);
+}
+
+template <>
+float AdditiveQuantizer::
+        compute_1_distance_LUT<false, AdditiveQuantizer::ST_LUT_nonorm>(
+                const uint8_t* codes,
+                const float* LUT) const {
+    BitstringReader bs(codes, code_size);
+    return -accumulate_IPs(*this, bs, codes, LUT);
+}
+
+template <>
+float AdditiveQuantizer::
+        compute_1_distance_LUT<false, AdditiveQuantizer::ST_norm_float>(
+                const uint8_t* codes,
+                const float* LUT) const {
+    BitstringReader bs(codes, code_size);
+    float accu = accumulate_IPs(*this, bs, codes, LUT);
+    uint32_t norm_i = bs.read(32);
+    float norm2;
+    memcpy(&norm2, &norm_i, 4);
+    return norm2 - 2 * accu;
+}
+
+template <>
+float AdditiveQuantizer::
+        compute_1_distance_LUT<false, AdditiveQuantizer::ST_norm_cqint8>(
+                const uint8_t* codes,
+                const float* LUT) const {
+    BitstringReader bs(codes, code_size);
+    float accu = accumulate_IPs(*this, bs, codes, LUT);
+    uint32_t norm_i = bs.read(8);
+    float norm2 = decode_qcint(norm_i);
+    return norm2 - 2 * accu;
+}
+
+template <>
+float AdditiveQuantizer::
+        compute_1_distance_LUT<false, AdditiveQuantizer::ST_norm_cqint4>(
+                const uint8_t* codes,
+                const float* LUT) const {
+    BitstringReader bs(codes, code_size);
+    float accu = accumulate_IPs(*this, bs, codes, LUT);
+    uint32_t norm_i = bs.read(4);
+    float norm2 = decode_qcint(norm_i);
+    return norm2 - 2 * accu;
+}
+
+template <>
+float AdditiveQuantizer::
+        compute_1_distance_LUT<false, AdditiveQuantizer::ST_norm_qint8>(
+                const uint8_t* codes,
+                const float* LUT) const {
+    BitstringReader bs(codes, code_size);
+    float accu = accumulate_IPs(*this, bs, codes, LUT);
+    uint32_t norm_i = bs.read(8);
+    float norm2 = decode_qint8(norm_i, norm_min, norm_max);
+    return norm2 - 2 * accu;
+}
+
+template <>
+float AdditiveQuantizer::
+        compute_1_distance_LUT<false, AdditiveQuantizer::ST_norm_qint4>(
+                const uint8_t* codes,
+                const float* LUT) const {
+    BitstringReader bs(codes, code_size);
+    float accu = accumulate_IPs(*this, bs, codes, LUT);
+    uint32_t norm_i = bs.read(4);
+    float norm2 = decode_qint4(norm_i, norm_min, norm_max);
+    return norm2 - 2 * accu;
+}
+
 } // namespace faiss

data/vendor/faiss/faiss/impl/AdditiveQuantizer.h

@@ -11,6 +11,8 @@
 #include <vector>
 
 #include <faiss/Index.h>
+#include <faiss/IndexFlat.h>
+#include <faiss/impl/Quantizer.h>
 
 namespace faiss {
 
@@ -20,58 +22,140 @@ namespace faiss {
  * concatenation of M sub-vectors, additive quantizers sum M sub-vectors
  * to get the decoded vector.
  */
-struct AdditiveQuantizer {
-    size_t d;                     ///< size of the input vectors
+struct AdditiveQuantizer : Quantizer {
     size_t M;                     ///< number of codebooks
     std::vector<size_t> nbits;    ///< bits for each step
     std::vector<float> codebooks; ///< codebooks
 
     // derived values
-    std::vector<size_t> codebook_offsets;
-    size_t code_size;           ///< code size in bytes
-    size_t tot_bits;            ///< total number of bits
+    std::vector<uint64_t> codebook_offsets;
+    size_t tot_bits;            ///< total number of bits (indexes + norms)
+    size_t norm_bits;           ///< bits allocated for the norms
     size_t total_codebook_size; ///< size of the codebook in vectors
-    bool is_byte_aligned;
+    bool only_8bit;             ///< are all nbits = 8 (use faster decoder)
 
     bool verbose;    ///< verbose during training?
     bool is_trained; ///< is trained or not
 
+    IndexFlat1D qnorm;            ///< store and search norms
+    std::vector<float> norm_tabs; ///< store norms of codebook entries for 4-bit
+                                  ///< fastscan search
+
+    /// norms and distance matrixes with beam search can get large, so use this
+    /// to control for the amount of memory that can be allocated
+    size_t max_mem_distances;
+
+    /// encode a norm into norm_bits bits
+    uint64_t encode_norm(float norm) const;
+
+    uint32_t encode_qcint(
+            float x) const; ///< encode norm by non-uniform scalar quantization
+
+    float decode_qcint(uint32_t c)
+            const; ///< decode norm by non-uniform scalar quantization
+
+    /// Encodes how search is performed and how vectors are encoded
+    enum Search_type_t {
+        ST_decompress,    ///< decompress database vector
+        ST_LUT_nonorm,    ///< use a LUT, don't include norms (OK for IP or
+                          ///< normalized vectors)
+        ST_norm_from_LUT, ///< compute the norms from the look-up tables (cost
+                          ///< is in O(M^2))
+        ST_norm_float, ///< use a LUT, and store float32 norm with the vectors
+        ST_norm_qint8, ///< use a LUT, and store 8bit-quantized norm
+        ST_norm_qint4,
+        ST_norm_cqint8, ///< use a LUT, and store non-uniform quantized norm
+        ST_norm_cqint4,
+
+        ST_norm_lsq2x4, ///< use a 2x4 bits lsq as norm quantizer (for fast
+                        ///< scan)
+        ST_norm_rq2x4,  ///< use a 2x4 bits rq as norm quantizer (for fast scan)
+    };
+
+    AdditiveQuantizer(
+            size_t d,
+            const std::vector<size_t>& nbits,
+            Search_type_t search_type = ST_decompress);
+
+    AdditiveQuantizer();
+
     ///< compute derived values when d, M and nbits have been set
     void set_derived_values();
 
-    ///< Train the additive quantizer
-    virtual void train(size_t n, const float* x) = 0;
+    ///< Train the norm quantizer
+    void train_norm(size_t n, const float* norms);
+
+    void compute_codes(const float* x, uint8_t* codes, size_t n)
+            const override {
+        compute_codes_add_centroids(x, codes, n);
+    }
 
     /** Encode a set of vectors
      *
      * @param x      vectors to encode, size n * d
      * @param codes  output codes, size n * code_size
+     * @param centroids  centroids to be added to x, size n * d
      */
-    virtual void compute_codes(const float* x, uint8_t* codes, size_t n)
-            const = 0;
+    virtual void compute_codes_add_centroids(
+            const float* x,
+            uint8_t* codes,
+            size_t n,
+            const float* centroids = nullptr) const = 0;
 
     /** pack a series of code to bit-compact format
      *
-     * @param codes  codes to be packed, size n * code_size
+     * @param codes        codes to be packed, size n * code_size
      * @param packed_codes output bit-compact codes
-     * @param ld_codes  leading dimension of codes
+     * @param ld_codes     leading dimension of codes
+     * @param norms        norms of the vectors (size n). Will be computed if
+     *                     needed but not provided
+     * @param centroids    centroids to be added to x, size n * d
      */
     void pack_codes(
             size_t n,
             const int32_t* codes,
             uint8_t* packed_codes,
-            int64_t ld_codes = -1) const;
+            int64_t ld_codes = -1,
+            const float* norms = nullptr,
+            const float* centroids = nullptr) const;
 
     /** Decode a set of vectors
      *
      * @param codes  codes to decode, size n * code_size
      * @param x      output vectors, size n * d
      */
-    void decode(const uint8_t* codes, float* x, size_t n) const;
+    void decode(const uint8_t* codes, float* x, size_t n) const override;
+
+    /** Decode a set of vectors in non-packed format
+     *
+     * @param codes  codes to decode, size n * ld_codes
+     * @param x      output vectors, size n * d
+     */
+    virtual void decode_unpacked(
+            const int32_t* codes,
+            float* x,
+            size_t n,
+            int64_t ld_codes = -1) const;
 
     /****************************************************************************
-     * Support for exhaustive distance computations with the centroids.
-     * Hence, the number of elements that can be enumerated is not too large.
+     * Search functions in an external set of codes.
+     ****************************************************************************/
+
+    /// Also determines what's in the codes
+    Search_type_t search_type;
+
+    /// min/max for quantization of norms
+    float norm_min, norm_max;
+
+    template <bool is_IP, Search_type_t effective_search_type>
+    float compute_1_distance_LUT(const uint8_t* codes, const float* LUT) const;
+
+    /*
+        float compute_1_L2sqr(const uint8_t* codes, const float* LUT);
+    */
+    /****************************************************************************
+     * Support for exhaustive distance computations with all the centroids.
+     * Hence, the number of these centroids should not be too large.
      ****************************************************************************/
     using idx_t = Index::idx_t;
 
@@ -83,11 +167,18 @@ struct AdditiveQuantizer {
      *
      * @param xq     query vector, size (n, d)
      * @param LUT    look-up table, size (n, total_codebook_size)
+     * @param alpha  compute alpha * inner-product
+     * @param ld_lut  leading dimension of LUT
      */
-    void compute_LUT(size_t n, const float* xq, float* LUT) const;
+    virtual void compute_LUT(
+            size_t n,
+            const float* xq,
+            float* LUT,
+            float alpha = 1.0f,
+            long ld_lut = -1) const;
 
     /// exact IP search
-    void knn_exact_inner_product(
+    void knn_centroids_inner_product(
             idx_t n,
             const float* xq,
             idx_t k,
@@ -101,7 +192,7 @@ struct AdditiveQuantizer {
     void compute_centroid_norms(float* norms) const;
 
     /** Exact L2 search, with precomputed norms */
-    void knn_exact_L2(
+    void knn_centroids_L2(
             idx_t n,
             const float* xq,
             idx_t k,