faiss 0.3.1 → 0.3.3

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

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.
@@ -24,7 +24,6 @@
 #include <faiss/utils/Heap.h>
 #include <faiss/utils/distances.h>
 #include <faiss/utils/hamming.h>
-#include <faiss/utils/utils.h>
 
 extern "C" {
 
@@ -152,6 +151,40 @@ void AdditiveQuantizer::train_norm(size_t n, const float* norms) {
     }
 }
 
+void AdditiveQuantizer::compute_codebook_tables() {
+    centroid_norms.resize(total_codebook_size);
+    fvec_norms_L2sqr(
+            centroid_norms.data(), codebooks.data(), d, total_codebook_size);
+    size_t cross_table_size = 0;
+    for (int m = 0; m < M; m++) {
+        size_t K = (size_t)1 << nbits[m];
+        cross_table_size += K * codebook_offsets[m];
+    }
+    codebook_cross_products.resize(cross_table_size);
+    size_t ofs = 0;
+    for (int m = 1; m < M; m++) {
+        FINTEGER ki = (size_t)1 << nbits[m];
+        FINTEGER kk = codebook_offsets[m];
+        FINTEGER di = d;
+        float zero = 0, one = 1;
+        assert(ofs + ki * kk <= cross_table_size);
+        sgemm_("Transposed",
+               "Not transposed",
+               &ki,
+               &kk,
+               &di,
+               &one,
+               codebooks.data() + d * kk,
+               &di,
+               codebooks.data(),
+               &di,
+               &zero,
+               codebook_cross_products.data() + ofs,
+               &ki);
+        ofs += ki * kk;
+    }
+}
+
 namespace {
 
 // TODO
@@ -471,7 +504,6 @@ 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++) {
@@ -483,6 +515,29 @@ float accumulate_IPs(
     return accu;
 }
 
+float compute_norm_from_LUT(const AdditiveQuantizer& aq, BitstringReader& bs) {
+    float accu = 0;
+    std::vector<int> idx(aq.M);
+    const float* c = aq.codebook_cross_products.data();
+    for (int m = 0; m < aq.M; m++) {
+        size_t nbit = aq.nbits[m];
+        int i = bs.read(nbit);
+        size_t K = 1 << nbit;
+        idx[m] = i;
+
+        accu += aq.centroid_norms[aq.codebook_offsets[m] + i];
+
+        for (int l = 0; l < m; l++) {
+            int j = idx[l];
+            accu += 2 * c[j * K + i];
+            c += (1 << aq.nbits[l]) * K;
+        }
+    }
+    // FAISS_THROW_IF_NOT(c == aq.codebook_cross_products.data() +
+    // aq.codebook_cross_products.size());
+    return accu;
+}
+
 } // anonymous namespace
 
 template <>
@@ -491,7 +546,7 @@ float AdditiveQuantizer::
                 const uint8_t* codes,
                 const float* LUT) const {
     BitstringReader bs(codes, code_size);
-    return accumulate_IPs(*this, bs, codes, LUT);
+    return accumulate_IPs(*this, bs, LUT);
 }
 
 template <>
@@ -500,7 +555,7 @@ float AdditiveQuantizer::
                 const uint8_t* codes,
                 const float* LUT) const {
     BitstringReader bs(codes, code_size);
-    return -accumulate_IPs(*this, bs, codes, LUT);
+    return -accumulate_IPs(*this, bs, LUT);
 }
 
 template <>
@@ -509,7 +564,7 @@ float AdditiveQuantizer::
                 const uint8_t* codes,
                 const float* LUT) const {
     BitstringReader bs(codes, code_size);
-    float accu = accumulate_IPs(*this, bs, codes, LUT);
+    float accu = accumulate_IPs(*this, bs, LUT);
     uint32_t norm_i = bs.read(32);
     float norm2;
     memcpy(&norm2, &norm_i, 4);
@@ -522,7 +577,7 @@ float AdditiveQuantizer::
                 const uint8_t* codes,
                 const float* LUT) const {
     BitstringReader bs(codes, code_size);
-    float accu = accumulate_IPs(*this, bs, codes, LUT);
+    float accu = accumulate_IPs(*this, bs, LUT);
     uint32_t norm_i = bs.read(8);
     float norm2 = decode_qcint(norm_i);
     return norm2 - 2 * accu;
@@ -534,7 +589,7 @@ float AdditiveQuantizer::
                 const uint8_t* codes,
                 const float* LUT) const {
     BitstringReader bs(codes, code_size);
-    float accu = accumulate_IPs(*this, bs, codes, LUT);
+    float accu = accumulate_IPs(*this, bs, LUT);
     uint32_t norm_i = bs.read(4);
     float norm2 = decode_qcint(norm_i);
     return norm2 - 2 * accu;
@@ -546,7 +601,7 @@ float AdditiveQuantizer::
                 const uint8_t* codes,
                 const float* LUT) const {
     BitstringReader bs(codes, code_size);
-    float accu = accumulate_IPs(*this, bs, codes, LUT);
+    float accu = accumulate_IPs(*this, bs, LUT);
     uint32_t norm_i = bs.read(8);
     float norm2 = decode_qint8(norm_i, norm_min, norm_max);
     return norm2 - 2 * accu;
@@ -558,10 +613,23 @@ float AdditiveQuantizer::
                 const uint8_t* codes,
                 const float* LUT) const {
     BitstringReader bs(codes, code_size);
-    float accu = accumulate_IPs(*this, bs, codes, LUT);
+    float accu = accumulate_IPs(*this, bs, LUT);
     uint32_t norm_i = bs.read(4);
     float norm2 = decode_qint4(norm_i, norm_min, norm_max);
     return norm2 - 2 * accu;
 }
 
+template <>
+float AdditiveQuantizer::
+        compute_1_distance_LUT<false, AdditiveQuantizer::ST_norm_from_LUT>(
+                const uint8_t* codes,
+                const float* LUT) const {
+    FAISS_THROW_IF_NOT(codebook_cross_products.size() > 0);
+    BitstringReader bs(codes, code_size);
+    float accu = accumulate_IPs(*this, bs, LUT);
+    BitstringReader bs2(codes, code_size);
+    float norm2 = compute_norm_from_LUT(*this, bs2);
+    return norm2 - 2 * accu;
+}
+
 } // namespace faiss

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

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.
@@ -29,6 +29,8 @@ struct AdditiveQuantizer : Quantizer {
     std::vector<float> codebooks; ///< codebooks
 
     // derived values
+    /// codebook #1 is stored in rows codebook_offsets[i]:codebook_offsets[i+1]
+    /// in the codebooks table of size total_codebook_size by d
     std::vector<uint64_t> codebook_offsets;
     size_t tot_bits = 0;            ///< total number of bits (indexes + norms)
     size_t norm_bits = 0;           ///< bits allocated for the norms
@@ -38,9 +40,19 @@ struct AdditiveQuantizer : Quantizer {
     bool verbose = false;    ///< verbose during training?
     bool is_trained = false; ///< 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
+    /// auxiliary data for ST_norm_lsq2x4 and ST_norm_rq2x4
+    /// store norms of codebook entries for 4-bit fastscan
+    std::vector<float> norm_tabs;
+    IndexFlat1D qnorm; ///< store and search norms
+
+    void compute_codebook_tables();
+
+    /// norms of all codebook entries (size total_codebook_size)
+    std::vector<float> centroid_norms;
+
+    /// dot products of all codebook entries with the previous codebooks
+    /// size sum(codebook_offsets[m] * 2^nbits[m], m=0..M-1)
+    std::vector<float> codebook_cross_products;
 
     /// norms and distance matrixes with beam search can get large, so use this
     /// to control for the amount of memory that can be allocated

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

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.
@@ -236,4 +236,29 @@ size_t InterruptCallback::get_period_hint(size_t flops) {
     return std::max((size_t)10 * 10 * 1000 * 1000 / (flops + 1), (size_t)1);
 }
 
+void TimeoutCallback::set_timeout(double timeout_in_seconds) {
+    timeout = timeout_in_seconds;
+    start = std::chrono::steady_clock::now();
+}
+
+bool TimeoutCallback::want_interrupt() {
+    if (timeout == 0) {
+        return false;
+    }
+    auto end = std::chrono::steady_clock::now();
+    std::chrono::duration<float, std::milli> duration = end - start;
+    float elapsed_in_seconds = duration.count() / 1000.0;
+    if (elapsed_in_seconds > timeout) {
+        timeout = 0;
+        return true;
+    }
+    return false;
+}
+
+void TimeoutCallback::reset(double timeout_in_seconds) {
+    auto tc(new faiss::TimeoutCallback());
+    faiss::InterruptCallback::instance.reset(tc);
+    tc->set_timeout(timeout_in_seconds);
+}
+
 } // namespace faiss

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

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.
@@ -122,7 +122,7 @@ struct RangeSearchPartialResult : BufferList {
     void copy_result(bool incremental = false);
 
     /// merge a set of PartialResult's into one RangeSearchResult
-    /// on ouptut the partialresults are empty!
+    /// on output the partialresults are empty!
     static void merge(
             std::vector<RangeSearchPartialResult*>& partial_results,
             bool do_delete = true);
@@ -161,6 +161,14 @@ struct FAISS_API InterruptCallback {
     static size_t get_period_hint(size_t flops);
 };
 
+struct TimeoutCallback : InterruptCallback {
+    std::chrono::time_point<std::chrono::steady_clock> start;
+    double timeout;
+    bool want_interrupt() override;
+    void set_timeout(double timeout_in_seconds);
+    static void reset(double timeout_in_seconds);
+};
+
 /// set implementation optimized for fast access.
 struct VisitedTable {
     std::vector<uint8_t> visited;

data/vendor/faiss/faiss/impl/CodePacker.cpp

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.

data/vendor/faiss/faiss/impl/CodePacker.h

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.

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

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.
@@ -59,6 +59,52 @@ struct DistanceComputer {
     virtual ~DistanceComputer() {}
 };
 
+/* Wrap the distance computer into one that negates the
+   distances. This makes supporting INNER_PRODUCE search easier */
+
+struct NegativeDistanceComputer : DistanceComputer {
+    /// owned by this
+    DistanceComputer* basedis;
+
+    explicit NegativeDistanceComputer(DistanceComputer* basedis)
+            : basedis(basedis) {}
+
+    void set_query(const float* x) override {
+        basedis->set_query(x);
+    }
+
+    /// compute distance of vector i to current query
+    float operator()(idx_t i) override {
+        return -(*basedis)(i);
+    }
+
+    void distances_batch_4(
+            const idx_t idx0,
+            const idx_t idx1,
+            const idx_t idx2,
+            const idx_t idx3,
+            float& dis0,
+            float& dis1,
+            float& dis2,
+            float& dis3) override {
+        basedis->distances_batch_4(
+                idx0, idx1, idx2, idx3, dis0, dis1, dis2, dis3);
+        dis0 = -dis0;
+        dis1 = -dis1;
+        dis2 = -dis2;
+        dis3 = -dis3;
+    }
+
+    /// compute distance between two stored vectors
+    float symmetric_dis(idx_t i, idx_t j) override {
+        return -basedis->symmetric_dis(i, j);
+    }
+
+    virtual ~NegativeDistanceComputer() {
+        delete basedis;
+    }
+};
+
 /*************************************************************
  * Specialized version of the DistanceComputer when we know that codes are
  * laid out in a flat index.

data/vendor/faiss/faiss/impl/FaissAssert.h

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.
@@ -94,13 +94,15 @@
         }                                              \
     } while (false)
 
-#define FAISS_THROW_IF_NOT_MSG(X, MSG)                       \
+#define FAISS_THROW_IF_MSG(X, MSG)                           \
     do {                                                     \
-        if (!(X)) {                                          \
+        if (X) {                                             \
             FAISS_THROW_FMT("Error: '%s' failed: " MSG, #X); \
         }                                                    \
     } while (false)
 
+#define FAISS_THROW_IF_NOT_MSG(X, MSG) FAISS_THROW_IF_MSG(!(X), MSG)
+
 #define FAISS_THROW_IF_NOT_FMT(X, FMT, ...)                               \
     do {                                                                  \
         if (!(X)) {                                                       \

data/vendor/faiss/faiss/impl/FaissException.cpp

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.

data/vendor/faiss/faiss/impl/FaissException.h

@@ -1,6 +1,5 @@
-
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.