faiss 0.2.0 → 0.2.4

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

@@ -0,0 +1,175 @@
+/**
+ * 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 <cstdint>
+#include <vector>
+
+#include <faiss/Index.h>
+#include <faiss/IndexFlat.h>
+
+namespace faiss {
+
+/** Abstract structure for additive quantizers
+ *
+ * Different from the product quantizer in which the decoded vector is the
+ * 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
+    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 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
+
+    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,
+    };
+
+    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;
+
+    /** Encode a set of vectors
+     *
+     * @param x      vectors to encode, 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;
+
+    /** pack a series of code to bit-compact format
+     *
+     * @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 norms        norms of the vectors (size n). Will be computed if
+     *                     needed but not provided
+     */
+    void pack_codes(
+            size_t n,
+            const int32_t* codes,
+            uint8_t* packed_codes,
+            int64_t ld_codes = -1,
+            const float* norms = 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;
+
+    /** 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(
+            const int32_t* codes,
+            float* x,
+            size_t n,
+            int64_t ld_codes = -1) const;
+
+    /****************************************************************************
+     * 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;
+
+    /// decoding function for a code in a 64-bit word
+    void decode_64bit(idx_t n, float* x) const;
+
+    /** Compute inner-product look-up tables. Used in the centroid search
+     * functions.
+     *
+     * @param xq     query vector, size (n, d)
+     * @param LUT    look-up table, size (n, total_codebook_size)
+     */
+    void compute_LUT(size_t n, const float* xq, float* LUT) const;
+
+    /// exact IP search
+    void knn_centroids_inner_product(
+            idx_t n,
+            const float* xq,
+            idx_t k,
+            float* distances,
+            idx_t* labels) const;
+
+    /** For L2 search we need the L2 norms of the centroids
+     *
+     * @param norms    output norms table, size total_codebook_size
+     */
+    void compute_centroid_norms(float* norms) const;
+
+    /** Exact L2 search, with precomputed norms */
+    void knn_centroids_L2(
+            idx_t n,
+            const float* xq,
+            idx_t k,
+            float* distances,
+            idx_t* labels,
+            const float* centroid_norms) const;
+
+    virtual ~AdditiveQuantizer();
+};
+
+}; // namespace faiss

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

@@ -14,18 +14,16 @@
 
 #include <faiss/impl/FaissAssert.h>
 
-
 namespace faiss {
 
-
 /***********************************************************************
  * RangeSearchResult
  ***********************************************************************/
 
-RangeSearchResult::RangeSearchResult (idx_t nq, bool alloc_lims): nq (nq) {
+RangeSearchResult::RangeSearchResult(idx_t nq, bool alloc_lims) : nq(nq) {
     if (alloc_lims) {
-        lims = new size_t [nq + 1];
-        memset (lims, 0, sizeof(*lims) * (nq + 1));
+        lims = new size_t[nq + 1];
+        memset(lims, 0, sizeof(*lims) * (nq + 1));
     } else {
         lims = nullptr;
     }
@@ -36,145 +34,129 @@ RangeSearchResult::RangeSearchResult (idx_t nq, bool alloc_lims): nq (nq) {
 
 /// called when lims contains the nb of elements result entries
 /// for each query
-void RangeSearchResult::do_allocation () {
+void RangeSearchResult::do_allocation() {
+    // works only if all the partial results are aggregated
+    // simulatenously
+    FAISS_THROW_IF_NOT(labels == nullptr && distances == nullptr);
     size_t ofs = 0;
     for (int i = 0; i < nq; i++) {
         size_t n = lims[i];
-        lims [i] = ofs;
+        lims[i] = ofs;
         ofs += n;
     }
-    lims [nq] = ofs;
-    labels = new idx_t [ofs];
-    distances = new float [ofs];
+    lims[nq] = ofs;
+    labels = new idx_t[ofs];
+    distances = new float[ofs];
 }
 
-RangeSearchResult::~RangeSearchResult () {
-    delete [] labels;
-    delete [] distances;
-    delete [] lims;
+RangeSearchResult::~RangeSearchResult() {
+    delete[] labels;
+    delete[] distances;
+    delete[] lims;
 }
 
-
-
-
-
 /***********************************************************************
  * BufferList
  ***********************************************************************/
 
-
-BufferList::BufferList (size_t buffer_size):
-    buffer_size (buffer_size)
-{
+BufferList::BufferList(size_t buffer_size) : buffer_size(buffer_size) {
     wp = buffer_size;
 }
 
-BufferList::~BufferList ()
-{
+BufferList::~BufferList() {
     for (int i = 0; i < buffers.size(); i++) {
-        delete [] buffers[i].ids;
-        delete [] buffers[i].dis;
+        delete[] buffers[i].ids;
+        delete[] buffers[i].dis;
     }
 }
 
-void BufferList::add (idx_t id, float dis) {
+void BufferList::add(idx_t id, float dis) {
     if (wp == buffer_size) { // need new buffer
         append_buffer();
     }
-    Buffer & buf = buffers.back();
-    buf.ids [wp] = id;
-    buf.dis [wp] = dis;
+    Buffer& buf = buffers.back();
+    buf.ids[wp] = id;
+    buf.dis[wp] = dis;
     wp++;
 }
 
-
-void BufferList::append_buffer ()
-{
-    Buffer buf = {new idx_t [buffer_size], new float [buffer_size]};
-    buffers.push_back (buf);
+void BufferList::append_buffer() {
+    Buffer buf = {new idx_t[buffer_size], new float[buffer_size]};
+    buffers.push_back(buf);
     wp = 0;
 }
 
 /// copy elemnts ofs:ofs+n-1 seen as linear data in the buffers to
 /// tables dest_ids, dest_dis
-void BufferList::copy_range (size_t ofs, size_t n,
-                             idx_t * dest_ids, float *dest_dis)
-{
+void BufferList::copy_range(
+        size_t ofs,
+        size_t n,
+        idx_t* dest_ids,
+        float* dest_dis) {
     size_t bno = ofs / buffer_size;
     ofs -= bno * buffer_size;
     while (n > 0) {
         size_t ncopy = ofs + n < buffer_size ? n : buffer_size - ofs;
-        Buffer buf = buffers [bno];
-        memcpy (dest_ids, buf.ids + ofs, ncopy * sizeof(*dest_ids));
-        memcpy (dest_dis, buf.dis + ofs, ncopy * sizeof(*dest_dis));
+        Buffer buf = buffers[bno];
+        memcpy(dest_ids, buf.ids + ofs, ncopy * sizeof(*dest_ids));
+        memcpy(dest_dis, buf.dis + ofs, ncopy * sizeof(*dest_dis));
         dest_ids += ncopy;
         dest_dis += ncopy;
         ofs = 0;
-        bno ++;
+        bno++;
         n -= ncopy;
     }
 }
 
-
 /***********************************************************************
  * RangeSearchPartialResult
  ***********************************************************************/
 
-void RangeQueryResult::add (float dis, idx_t id) {
+void RangeQueryResult::add(float dis, idx_t id) {
     nres++;
-    pres->add (id, dis);
+    pres->add(id, dis);
 }
 
-
-
-RangeSearchPartialResult::RangeSearchPartialResult (RangeSearchResult * res_in):
-    BufferList(res_in->buffer_size),
-    res(res_in)
-{}
-
+RangeSearchPartialResult::RangeSearchPartialResult(RangeSearchResult* res_in)
+        : BufferList(res_in->buffer_size), res(res_in) {}
 
 /// begin a new result
-RangeQueryResult &
-    RangeSearchPartialResult::new_result (idx_t qno)
-{
+RangeQueryResult& RangeSearchPartialResult::new_result(idx_t qno) {
     RangeQueryResult qres = {qno, 0, this};
-    queries.push_back (qres);
+    queries.push_back(qres);
     return queries.back();
 }
 
-
-void RangeSearchPartialResult::finalize ()
-{
-    set_lims ();
+void RangeSearchPartialResult::finalize() {
+    set_lims();
 #pragma omp barrier
 
 #pragma omp single
-    res->do_allocation ();
+    res->do_allocation();
 
 #pragma omp barrier
-    copy_result ();
+    copy_result();
 }
 
-
 /// called by range_search before do_allocation
-void RangeSearchPartialResult::set_lims ()
-{
+void RangeSearchPartialResult::set_lims() {
     for (int i = 0; i < queries.size(); i++) {
-        RangeQueryResult & qres = queries[i];
+        RangeQueryResult& qres = queries[i];
         res->lims[qres.qno] = qres.nres;
     }
 }
 
 /// called by range_search after do_allocation
-void RangeSearchPartialResult::copy_result (bool incremental)
-{
+void RangeSearchPartialResult::copy_result(bool incremental) {
     size_t ofs = 0;
     for (int i = 0; i < queries.size(); i++) {
-        RangeQueryResult & qres = queries[i];
+        RangeQueryResult& qres = queries[i];
 
-        copy_range (ofs, qres.nres,
-                    res->labels + res->lims[qres.qno],
-                    res->distances + res->lims[qres.qno]);
+        copy_range(
+                ofs,
+                qres.nres,
+                res->labels + res->lims[qres.qno],
+                res->distances + res->lims[qres.qno]);
         if (incremental) {
             res->lims[qres.qno] += qres.nres;
         }
@@ -182,26 +164,28 @@ void RangeSearchPartialResult::copy_result (bool incremental)
     }
 }
 
-void RangeSearchPartialResult::merge (std::vector <RangeSearchPartialResult *> &
-                                      partial_results, bool do_delete)
-{
-
+void RangeSearchPartialResult::merge(
+        std::vector<RangeSearchPartialResult*>& partial_results,
+        bool do_delete) {
     int npres = partial_results.size();
-    if (npres == 0) return;
-    RangeSearchResult *result = partial_results[0]->res;
+    if (npres == 0)
+        return;
+    RangeSearchResult* result = partial_results[0]->res;
     size_t nx = result->nq;
 
     // count
-    for (const RangeSearchPartialResult * pres : partial_results) {
-        if (!pres) continue;
-        for (const RangeQueryResult &qres : pres->queries) {
+    for (const RangeSearchPartialResult* pres : partial_results) {
+        if (!pres)
+            continue;
+        for (const RangeQueryResult& qres : pres->queries) {
             result->lims[qres.qno] += qres.nres;
         }
     }
-    result->do_allocation ();
+    result->do_allocation();
     for (int j = 0; j < npres; j++) {
-        if (!partial_results[j]) continue;
-        partial_results[j]->copy_result (true);
+        if (!partial_results[j])
+            continue;
+        partial_results[j]->copy_result(true);
         if (do_delete) {
             delete partial_results[j];
             partial_results[j] = nullptr;
@@ -210,22 +194,19 @@ void RangeSearchPartialResult::merge (std::vector <RangeSearchPartialResult *> &
 
     // reset the limits
     for (size_t i = nx; i > 0; i--) {
-        result->lims [i] = result->lims [i - 1];
+        result->lims[i] = result->lims[i - 1];
     }
-    result->lims [0] = 0;
+    result->lims[0] = 0;
 }
 
 /***********************************************************************
  * IDSelectorRange
  ***********************************************************************/
 
-IDSelectorRange::IDSelectorRange (idx_t imin, idx_t imax):
-    imin (imin), imax (imax)
-{
-}
+IDSelectorRange::IDSelectorRange(idx_t imin, idx_t imax)
+        : imin(imin), imax(imax) {}
 
-bool IDSelectorRange::is_member (idx_t id) const
-{
+bool IDSelectorRange::is_member(idx_t id) const {
     return id >= imin && id < imax;
 }
 
@@ -233,33 +214,29 @@ bool IDSelectorRange::is_member (idx_t id) const
  * IDSelectorArray
  ***********************************************************************/
 
-IDSelectorArray::IDSelectorArray (size_t n, const idx_t *ids):
-   n (n), ids(ids)
-{
-}
+IDSelectorArray::IDSelectorArray(size_t n, const idx_t* ids) : n(n), ids(ids) {}
 
-bool IDSelectorArray::is_member (idx_t id) const
-{
+bool IDSelectorArray::is_member(idx_t id) const {
     for (idx_t i = 0; i < n; i++) {
-        if (ids[i] == id) return true;
+        if (ids[i] == id)
+            return true;
     }
     return false;
 }
 
-
 /***********************************************************************
  * IDSelectorBatch
  ***********************************************************************/
 
-IDSelectorBatch::IDSelectorBatch (size_t n, const idx_t *indices)
-{
+IDSelectorBatch::IDSelectorBatch(size_t n, const idx_t* indices) {
     nbits = 0;
-    while (n > (1L << nbits)) nbits++;
+    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);
+    bloom.resize(1UL << (nbits - 3), 0);
     for (long i = 0; i < n; i++) {
         Index::idx_t id = indices[i];
         set.insert(id);
@@ -268,39 +245,36 @@ IDSelectorBatch::IDSelectorBatch (size_t n, const idx_t *indices)
     }
 }
 
-bool IDSelectorBatch::is_member (idx_t i) const
-{
+bool IDSelectorBatch::is_member(idx_t i) const {
     long im = i & mask;
-    if(!(bloom[im>>3] & (1 << (im & 7)))) {
+    if (!(bloom[im >> 3] & (1 << (im & 7)))) {
         return 0;
     }
     return set.count(i);
 }
 
-
 /***********************************************************
  * Interrupt callback
  ***********************************************************/
 
-
 std::unique_ptr<InterruptCallback> InterruptCallback::instance;
 
 std::mutex InterruptCallback::lock;
 
-void InterruptCallback::clear_instance () {
-    delete instance.release ();
+void InterruptCallback::clear_instance() {
+    delete instance.release();
 }
 
-void InterruptCallback::check () {
+void InterruptCallback::check() {
     if (!instance.get()) {
         return;
     }
-    if (instance->want_interrupt ()) {
-        FAISS_THROW_MSG ("computation interrupted");
+    if (instance->want_interrupt()) {
+        FAISS_THROW_MSG("computation interrupted");
     }
 }
 
-bool InterruptCallback::is_interrupted () {
+bool InterruptCallback::is_interrupted() {
     if (!instance.get()) {
         return false;
     }
@@ -308,8 +282,7 @@ bool InterruptCallback::is_interrupted () {
     return instance->want_interrupt();
 }
 
-
-size_t InterruptCallback::get_period_hint (size_t flops) {
+size_t InterruptCallback::get_period_hint(size_t flops) {
     if (!instance.get()) {
         return 1L << 30; // never check
     }
@@ -317,7 +290,4 @@ size_t InterruptCallback::get_period_hint (size_t flops) {
     return std::max((size_t)10 * 10 * 1000 * 1000 / (flops + 1), (size_t)1);
 }
 
-
-
-
 } // namespace faiss