faiss 0.2.4 → 0.2.5

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

@@ -17,7 +17,10 @@
 
 #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>
@@ -48,14 +51,14 @@ AdditiveQuantizer::AdditiveQuantizer(
         size_t d,
         const std::vector<size_t>& nbits,
         Search_type_t search_type)
-        : d(d),
+        : 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;
-    code_size = 0;
     tot_bits = 0;
     total_codebook_size = 0;
     only_8bit = false;
@@ -80,27 +83,82 @@ void AdditiveQuantizer::set_derived_values() {
     }
     total_codebook_size = codebook_offsets[M];
     switch (search_type) {
-        case ST_decompress:
-        case ST_LUT_nonorm:
-        case ST_norm_from_LUT:
-            break; // nothing to add
         case ST_norm_float:
-            tot_bits += 32;
+            norm_bits = 32;
             break;
         case ST_norm_qint8:
         case ST_norm_cqint8:
-            tot_bits += 8;
+        case ST_norm_lsq2x4:
+        case ST_norm_rq2x4:
+            norm_bits = 8;
             break;
         case ST_norm_qint4:
         case ST_norm_cqint4:
-            tot_bits += 4;
+            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
@@ -132,7 +190,7 @@ float decode_qint4(uint8_t i, float amin, float amax) {
 
 uint32_t AdditiveQuantizer::encode_qcint(float x) const {
     idx_t id;
-    qnorm.assign(idx_t(1), &x, &id, idx_t(1));
+    qnorm.assign(1, &x, &id, 1);
     return uint32_t(id);
 }
 
@@ -140,23 +198,54 @@ 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 float* norms) const {
+        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) {
-        if (!norms) {
+        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();
         }
@@ -168,34 +257,8 @@ void AdditiveQuantizer::pack_codes(
         for (int m = 0; m < M; m++) {
             bsw.write(codes1[m], nbits[m]);
         }
-        switch (search_type) {
-            case ST_decompress:
-            case ST_LUT_nonorm:
-            case ST_norm_from_LUT:
-                break;
-            case ST_norm_float:
-                bsw.write(*(uint32_t*)&norms[i], 32);
-                break;
-            case ST_norm_qint8: {
-                uint8_t b = encode_qint8(norms[i], norm_min, norm_max);
-                bsw.write(b, 8);
-                break;
-            }
-            case ST_norm_qint4: {
-                uint8_t b = encode_qint4(norms[i], norm_min, norm_max);
-                bsw.write(b, 4);
-                break;
-            }
-            case ST_norm_cqint8: {
-                uint32_t b = encode_qcint(norms[i]);
-                bsw.write(b, 8);
-                break;
-            }
-            case ST_norm_cqint4: {
-                uint32_t b = encode_qcint(norms[i]);
-                bsw.write(b, 4);
-                break;
-            }
+        if (norm_bits != 0) {
+            bsw.write(encode_norm(norms[i]), norm_bits);
         }
     }
 }
@@ -283,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 {
@@ -448,7 +516,8 @@ float AdditiveQuantizer::
     BitstringReader bs(codes, code_size);
     float accu = accumulate_IPs(*this, bs, codes, LUT);
     uint32_t norm_i = bs.read(32);
-    float norm2 = *(float*)&norm_i;
+    float norm2;
+    memcpy(&norm2, &norm_i, 4);
     return norm2 - 2 * accu;
 }
 

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

@@ -12,6 +12,7 @@
 
 #include <faiss/Index.h>
 #include <faiss/IndexFlat.h>
+#include <faiss/impl/Quantizer.h>
 
 namespace faiss {
 
@@ -21,23 +22,31 @@ 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<uint64_t> codebook_offsets;
-    size_t code_size;           ///< code size in bytes
-    size_t tot_bits;            ///< total number of bits
+    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 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
+    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
@@ -57,6 +66,10 @@ struct AdditiveQuantizer {
         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(
@@ -69,16 +82,25 @@ struct 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
      *
@@ -87,27 +109,29 @@ struct AdditiveQuantizer {
      * @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 float* norms = nullptr) const;
+            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
      */
-    void decode_unpacked(
+    virtual void decode_unpacked(
             const int32_t* codes,
             float* x,
             size_t n,
@@ -143,8 +167,15 @@ 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_centroids_inner_product(

data/vendor/faiss/faiss/impl/AuxIndexStructures.cpp

@@ -199,60 +199,6 @@ void RangeSearchPartialResult::merge(
     result->lims[0] = 0;
 }
 
-/***********************************************************************
- * IDSelectorRange
- ***********************************************************************/
-
-IDSelectorRange::IDSelectorRange(idx_t imin, idx_t imax)
-        : imin(imin), imax(imax) {}
-
-bool IDSelectorRange::is_member(idx_t id) const {
-    return id >= imin && id < imax;
-}
-
-/***********************************************************************
- * IDSelectorArray
- ***********************************************************************/
-
-IDSelectorArray::IDSelectorArray(size_t n, const idx_t* ids) : n(n), ids(ids) {}
-
-bool IDSelectorArray::is_member(idx_t id) const {
-    for (idx_t i = 0; i < n; i++) {
-        if (ids[i] == id)
-            return true;
-    }
-    return false;
-}
-
-/***********************************************************************
- * IDSelectorBatch
- ***********************************************************************/
-
-IDSelectorBatch::IDSelectorBatch(size_t n, const idx_t* indices) {
-    nbits = 0;
-    while (n > (1L << nbits))
-        nbits++;
-    nbits += 5;
-    // for n = 1M, nbits = 25 is optimal, see P56659518
-
-    mask = (1L << nbits) - 1;
-    bloom.resize(1UL << (nbits - 3), 0);
-    for (long i = 0; i < n; i++) {
-        Index::idx_t id = indices[i];
-        set.insert(id);
-        id &= mask;
-        bloom[id >> 3] |= 1 << (id & 7);
-    }
-}
-
-bool IDSelectorBatch::is_member(idx_t i) const {
-    long im = i & mask;
-    if (!(bloom[im >> 3] & (1 << (im & 7)))) {
-        return 0;
-    }
-    return set.count(i);
-}
-
 /***********************************************************
  * Interrupt callback
  ***********************************************************/

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

@@ -5,8 +5,6 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-// -*- c++ -*-
-
 // Auxiliary index structures, that are used in indexes but that can
 // be forward-declared
 
@@ -18,7 +16,6 @@
 #include <cstring>
 #include <memory>
 #include <mutex>
-#include <unordered_set>
 #include <vector>
 
 #include <faiss/Index.h>
@@ -52,55 +49,6 @@ struct RangeSearchResult {
     virtual ~RangeSearchResult();
 };
 
-/** Encapsulates a set of ids to remove. */
-struct IDSelector {
-    typedef Index::idx_t idx_t;
-    virtual bool is_member(idx_t id) const = 0;
-    virtual ~IDSelector() {}
-};
-
-/** remove ids between [imni, imax) */
-struct IDSelectorRange : IDSelector {
-    idx_t imin, imax;
-
-    IDSelectorRange(idx_t imin, idx_t imax);
-    bool is_member(idx_t id) const override;
-    ~IDSelectorRange() override {}
-};
-
-/** simple list of elements to remove
- *
- * this is inefficient in most cases, except for IndexIVF with
- * maintain_direct_map
- */
-struct IDSelectorArray : IDSelector {
-    size_t n;
-    const idx_t* ids;
-
-    IDSelectorArray(size_t n, const idx_t* ids);
-    bool is_member(idx_t id) const override;
-    ~IDSelectorArray() override {}
-};
-
-/** Remove ids from a set. Repetitions of ids in the indices set
- * passed to the constructor does not hurt performance. The hash
- * function used for the bloom filter and GCC's implementation of
- * unordered_set are just the least significant bits of the id. This
- * works fine for random ids or ids in sequences but will produce many
- * hash collisions if lsb's are always the same */
-struct IDSelectorBatch : IDSelector {
-    std::unordered_set<idx_t> set;
-
-    typedef unsigned char uint8_t;
-    std::vector<uint8_t> bloom; // assumes low bits of id are a good hash value
-    int nbits;
-    idx_t mask;
-
-    IDSelectorBatch(size_t n, const idx_t* indices);
-    bool is_member(idx_t id) const override;
-    ~IDSelectorBatch() override {}
-};
-
 /****************************************************************
  * Result structures for range search.
  *
@@ -186,30 +134,6 @@ struct RangeSearchPartialResult : BufferList {
             bool do_delete = true);
 };
 
-/***********************************************************
- * The distance computer maintains a current query and computes
- * distances to elements in an index that supports random access.
- *
- * The DistanceComputer is not intended to be thread-safe (eg. because
- * it maintains counters) so the distance functions are not const,
- * instantiate one from each thread if needed.
- ***********************************************************/
-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;
-
-    /// compute distance of vector i to current query
-    virtual float operator()(idx_t i) = 0;
-
-    /// compute distance between two stored vectors
-    virtual float symmetric_dis(idx_t i, idx_t j) = 0;
-
-    virtual ~DistanceComputer() {}
-};
-
 /***********************************************************
  * Interrupt callback
  ***********************************************************/

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

@@ -0,0 +1,64 @@
+/**
+ * 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/Index.h>
+
+namespace faiss {
+
+/***********************************************************
+ * The distance computer maintains a current query and computes
+ * distances to elements in an index that supports random access.
+ *
+ * The DistanceComputer is not intended to be thread-safe (eg. because
+ * it maintains counters) so the distance functions are not const,
+ * instantiate one from each thread if needed.
+ *
+ * Note that the equivalent for IVF indexes is the InvertedListScanner,
+ * 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;
+
+    /// compute distance of vector i to current query
+    virtual float operator()(idx_t i) = 0;
+
+    /// compute distance between two stored vectors
+    virtual float symmetric_dis(idx_t i, idx_t j) = 0;
+
+    virtual ~DistanceComputer() {}
+};
+
+/*************************************************************
+ * Specialized version of the DistanceComputer when we know that codes are
+ * laid out in a flat index.
+ */
+struct FlatCodesDistanceComputer : DistanceComputer {
+    const uint8_t* codes;
+    size_t code_size;
+
+    FlatCodesDistanceComputer(const uint8_t* codes, size_t code_size)
+            : codes(codes), code_size(code_size) {}
+
+    FlatCodesDistanceComputer() : codes(nullptr), code_size(0) {}
+
+    float operator()(idx_t i) final {
+        return distance_to_code(codes + i * code_size);
+    }
+
+    /// compute distance of current query to an encoded vector
+    virtual float distance_to_code(const uint8_t* code) = 0;
+
+    virtual ~FlatCodesDistanceComputer() {}
+};
+
+} // namespace faiss