faiss 0.2.3 → 0.2.5

data/vendor/faiss/faiss/impl/kmeans1d.cpp

@@ -0,0 +1,303 @@
+/**
+ * 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.
+ */
+
+#include <algorithm>
+#include <cstdint>
+#include <cstring>
+#include <functional>
+#include <numeric>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include <faiss/Index.h>
+#include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/kmeans1d.h>
+
+namespace faiss {
+
+using idx_t = Index::idx_t;
+using LookUpFunc = std::function<float(idx_t, idx_t)>;
+
+void reduce(
+        const std::vector<idx_t>& rows,
+        const std::vector<idx_t>& input_cols,
+        const LookUpFunc& lookup,
+        std::vector<idx_t>& output_cols) {
+    for (idx_t col : input_cols) {
+        while (!output_cols.empty()) {
+            idx_t row = rows[output_cols.size() - 1];
+            float a = lookup(row, col);
+            float b = lookup(row, output_cols.back());
+            if (a >= b) { // defeated
+                break;
+            }
+            output_cols.pop_back();
+        }
+        if (output_cols.size() < rows.size()) {
+            output_cols.push_back(col);
+        }
+    }
+}
+
+void interpolate(
+        const std::vector<idx_t>& rows,
+        const std::vector<idx_t>& cols,
+        const LookUpFunc& lookup,
+        idx_t* argmins) {
+    std::unordered_map<idx_t, idx_t> idx_to_col;
+    for (idx_t idx = 0; idx < cols.size(); ++idx) {
+        idx_to_col[cols[idx]] = idx;
+    }
+
+    idx_t start = 0;
+    for (idx_t r = 0; r < rows.size(); r += 2) {
+        idx_t row = rows[r];
+        idx_t end = cols.size() - 1;
+        if (r < rows.size() - 1) {
+            idx_t idx = argmins[rows[r + 1]];
+            end = idx_to_col[idx];
+        }
+        idx_t argmin = cols[start];
+        float min = lookup(row, argmin);
+        for (idx_t c = start + 1; c <= end; c++) {
+            float value = lookup(row, cols[c]);
+            if (value < min) {
+                argmin = cols[c];
+                min = value;
+            }
+        }
+        argmins[row] = argmin;
+        start = end;
+    }
+}
+
+/** SMAWK algo. Find the row minima of a monotone matrix.
+ *
+ * References:
+ *   1. http://web.cs.unlv.edu/larmore/Courses/CSC477/monge.pdf
+ *   2. https://gist.github.com/dstein64/8e94a6a25efc1335657e910ff525f405
+ *   3. https://github.com/dstein64/kmeans1d
+ */
+void smawk_impl(
+        const std::vector<idx_t>& rows,
+        const std::vector<idx_t>& input_cols,
+        const LookUpFunc& lookup,
+        idx_t* argmins) {
+    if (rows.size() == 0) {
+        return;
+    }
+
+    /**********************************
+     * REDUCE
+     **********************************/
+    auto ptr = &input_cols;
+    std::vector<idx_t> survived_cols; // survived columns
+    if (rows.size() < input_cols.size()) {
+        reduce(rows, input_cols, lookup, survived_cols);
+        ptr = &survived_cols;
+    }
+    auto& cols = *ptr; // avoid memory copy
+
+    /**********************************
+     * INTERPOLATE
+     **********************************/
+
+    // call recursively on odd-indexed rows
+    std::vector<idx_t> odd_rows;
+    for (idx_t i = 1; i < rows.size(); i += 2) {
+        odd_rows.push_back(rows[i]);
+    }
+    smawk_impl(odd_rows, cols, lookup, argmins);
+
+    // interpolate the even-indexed rows
+    interpolate(rows, cols, lookup, argmins);
+}
+
+void smawk(
+        const idx_t nrows,
+        const idx_t ncols,
+        const LookUpFunc& lookup,
+        idx_t* argmins) {
+    std::vector<idx_t> rows(nrows);
+    std::vector<idx_t> cols(ncols);
+    std::iota(std::begin(rows), std::end(rows), 0);
+    std::iota(std::begin(cols), std::end(cols), 0);
+
+    smawk_impl(rows, cols, lookup, argmins);
+}
+
+void smawk(
+        const idx_t nrows,
+        const idx_t ncols,
+        const float* x,
+        idx_t* argmins) {
+    auto lookup = [&x, &ncols](idx_t i, idx_t j) { return x[i * ncols + j]; };
+    smawk(nrows, ncols, lookup, argmins);
+}
+
+namespace {
+
+class CostCalculator {
+    // The reuslt would be inaccurate if we use float
+    std::vector<double> cumsum;
+    std::vector<double> cumsum2;
+
+   public:
+    CostCalculator(const std::vector<float>& vec, idx_t n) {
+        cumsum.push_back(0.0);
+        cumsum2.push_back(0.0);
+        for (idx_t i = 0; i < n; ++i) {
+            float x = vec[i];
+            cumsum.push_back(x + cumsum[i]);
+            cumsum2.push_back(x * x + cumsum2[i]);
+        }
+    }
+
+    float operator()(idx_t i, idx_t j) {
+        if (j < i) {
+            return 0.0f;
+        }
+        auto mu = (cumsum[j + 1] - cumsum[i]) / (j - i + 1);
+        auto result = cumsum2[j + 1] - cumsum2[i];
+        result += (j - i + 1) * (mu * mu);
+        result -= (2 * mu) * (cumsum[j + 1] - cumsum[i]);
+        return float(result);
+    }
+};
+
+template <class T>
+class Matrix {
+    std::vector<T> data;
+    idx_t nrows;
+    idx_t ncols;
+
+   public:
+    Matrix(idx_t nrows, idx_t ncols) {
+        this->nrows = nrows;
+        this->ncols = ncols;
+        data.resize(nrows * ncols);
+    }
+
+    inline T& at(idx_t i, idx_t j) {
+        return data[i * ncols + j];
+    }
+};
+
+} // anonymous namespace
+
+double kmeans1d(const float* x, size_t n, size_t nclusters, float* centroids) {
+    FAISS_THROW_IF_NOT(n >= nclusters);
+
+    // corner case
+    if (n == nclusters) {
+        memcpy(centroids, x, n * sizeof(*x));
+        return 0.0f;
+    }
+
+    /***************************************************
+     * sort in ascending order, O(NlogN) in time
+     ***************************************************/
+    std::vector<float> arr(x, x + n);
+    std::sort(arr.begin(), arr.end());
+
+    /***************************************************
+    dynamic programming algorithm
+
+    Reference: https://arxiv.org/abs/1701.07204
+    -------------------------------
+
+    Assume x is already sorted in ascending order.
+
+    N: number of points
+    K: number of clusters
+
+    CC(i, j): the cost of grouping xi,...,xj into one cluster
+    D[k][m]:  the cost of optimally clustering x1,...,xm into k clusters
+    T[k][m]:  the start index of the k-th cluster
+
+    The DP process is as follow:
+        D[k][m] = min_i D[k − 1][i − 1] + CC(i, m)
+        T[k][m] = argmin_i D[k − 1][i − 1] + CC(i, m)
+
+    This could be solved in O(KN^2) time and O(KN) space.
+
+    To further reduce the time complexity, we use SMAWK algo to
+    solve the argmin problem as follow:
+
+    For each k:
+        C[m][i] = D[k − 1][i − 1] + CC(i, m)
+
+        Here C is a n x n totally monotone matrix.
+        We could find the row minima by SMAWK in O(N) time.
+
+    Now the time complexity is reduced from O(kN^2) to O(KN).
+    ****************************************************/
+
+    CostCalculator CC(arr, n);
+    Matrix<float> D(nclusters, n);
+    Matrix<idx_t> T(nclusters, n);
+
+    for (idx_t m = 0; m < n; m++) {
+        D.at(0, m) = CC(0, m);
+        T.at(0, m) = 0;
+    }
+
+    std::vector<idx_t> indices(nclusters, 0);
+
+    for (idx_t k = 1; k < nclusters; ++k) {
+        // we define C here
+        auto C = [&D, &CC, &k](idx_t m, idx_t i) {
+            if (i == 0) {
+                return CC(i, m);
+            }
+            idx_t col = std::min(m, i - 1);
+            return D.at(k - 1, col) + CC(i, m);
+        };
+
+        std::vector<idx_t> argmins(n); // argmin of each row
+        smawk(n, n, C, argmins.data());
+        for (idx_t m = 0; m < argmins.size(); m++) {
+            idx_t idx = argmins[m];
+            D.at(k, m) = C(m, idx);
+            T.at(k, m) = idx;
+        }
+    }
+
+    /***************************************************
+    compute centroids by backtracking
+
+           T[K - 1][T[K][N] - 1]        T[K][N]        N
+    --------------|------------------------|-----------|
+                  |     cluster K - 1      | cluster K |
+
+    ****************************************************/
+
+    // for imbalance factor
+    double tot = 0.0;
+    double uf = 0.0;
+
+    idx_t end = n;
+    for (idx_t k = nclusters - 1; k >= 0; k--) {
+        const idx_t start = T.at(k, end - 1);
+        const float sum =
+                std::accumulate(arr.data() + start, arr.data() + end, 0.0f);
+        const idx_t size = end - start;
+        FAISS_THROW_IF_NOT_FMT(
+                size > 0, "Cluster %d: size %d", int(k), int(size));
+        centroids[k] = sum / size;
+        end = start;
+
+        tot += size;
+        uf += size * double(size);
+    }
+
+    uf = uf * nclusters / (tot * tot);
+    return uf;
+}
+
+} // namespace faiss

data/vendor/faiss/faiss/impl/kmeans1d.h

@@ -0,0 +1,48 @@
+/**
+ * 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>
+#include <functional>
+
+namespace faiss {
+
+/** SMAWK algorithm. Find the row minima of a monotone matrix.
+ *
+ * Expose this for testing.
+ *
+ * @param nrows    number of rows
+ * @param ncols    number of columns
+ * @param x        input matrix, size (nrows, ncols)
+ * @param argmins  argmin of each row
+ */
+void smawk(
+        const Index::idx_t nrows,
+        const Index::idx_t ncols,
+        const float* x,
+        Index::idx_t* argmins);
+
+/** Exact 1D K-Means by dynamic programming
+ *
+ * From  "Fast Exact k-Means, k-Medians and Bregman Divergence Clustering in 1D"
+ * Allan Grønlund, Kasper Green Larsen, Alexander Mathiasen, Jesper Sindahl
+ * Nielsen, Stefan Schneider, Mingzhou Song, ArXiV'17
+ *
+ * Section 2.2
+ *
+ * https://arxiv.org/abs/1701.07204
+ *
+ * @param x          input 1D array
+ * @param n          input array length
+ * @param nclusters  number of clusters
+ * @param centroids  output centroids, size nclusters
+ * @return  imbalancce factor
+ */
+double kmeans1d(const float* x, size_t n, size_t nclusters, float* centroids);
+
+} // namespace faiss

data/vendor/faiss/faiss/impl/pq4_fast_scan.cpp

@@ -122,30 +122,70 @@ void pq4_pack_codes_range(
     }
 }
 
+namespace {
+
+// get the specific address of the vector inside a block
+// shift is used for determine the if the saved in bits 0..3 (false) or
+// bits 4..7 (true)
+uint8_t get_vector_specific_address(
+        size_t bbs,
+        size_t vector_id,
+        size_t sq,
+        bool& shift) {
+    // get the vector_id inside the block
+    vector_id = vector_id % bbs;
+    shift = vector_id > 15;
+    vector_id = vector_id & 15;
+
+    // get the address of the vector in sq
+    size_t address;
+    if (vector_id < 8) {
+        address = vector_id << 1;
+    } else {
+        address = ((vector_id - 8) << 1) + 1;
+    }
+    if (sq & 1) {
+        address += 16;
+    }
+    return (sq >> 1) * bbs + address;
+}
+
+} // anonymous namespace
+
 uint8_t pq4_get_packed_element(
         const uint8_t* data,
         size_t bbs,
         size_t nsq,
-        size_t i,
+        size_t vector_id,
         size_t sq) {
     // move to correct bbs-sized block
-    data += (i / bbs * (nsq / 2) + sq / 2) * bbs;
-    sq = sq & 1;
-    i = i % bbs;
-
-    // another step
-    data += (i / 32) * 32;
-    i = i % 32;
-
-    if (sq == 1) {
-        data += 16;
+    // number of blocks * block size
+    data += (vector_id / bbs) * (((nsq + 1) / 2) * bbs);
+    bool shift;
+    size_t address = get_vector_specific_address(bbs, vector_id, sq, shift);
+    if (shift) {
+        return data[address] >> 4;
+    } else {
+        return data[address] & 15;
     }
-    const uint8_t iperm0[16] = {
-            0, 2, 4, 6, 8, 10, 12, 14, 1, 3, 5, 7, 9, 11, 13, 15};
-    if (i < 16) {
-        return data[iperm0[i]] & 15;
+}
+
+void pq4_set_packed_element(
+        uint8_t* data,
+        uint8_t code,
+        size_t bbs,
+        size_t nsq,
+        size_t vector_id,
+        size_t sq) {
+    // move to correct bbs-sized block
+    // number of blocks * block size
+    data += (vector_id / bbs) * (((nsq + 1) / 2) * bbs);
+    bool shift;
+    size_t address = get_vector_specific_address(bbs, vector_id, sq, shift);
+    if (shift) {
+        data[address] = (code << 4) | (data[address] & 15);
     } else {
-        return data[iperm0[i - 16]] >> 4;
+        data[address] = code | (data[address] & ~15);
     }
 }
 

data/vendor/faiss/faiss/impl/pq4_fast_scan.h

@@ -26,7 +26,7 @@ namespace faiss {
  *  The unused bytes are set to 0.
  *
  * @param codes   input codes, size (ntotal, ceil(M / 2))
- * @param nototal number of input codes
+ * @param ntotal  number of input codes
  * @param nb      output number of codes (ntotal rounded up to a multiple of
  *                bbs)
  * @param M2      number of sub-quantizers (=M rounded up to a muliple of 2)
@@ -61,14 +61,27 @@ void pq4_pack_codes_range(
 
 /** get a single element from a packed codes table
  *
- * @param i        vector id
+ * @param vector_id        vector id
  * @param sq       subquantizer (< nsq)
  */
 uint8_t pq4_get_packed_element(
         const uint8_t* data,
         size_t bbs,
         size_t nsq,
-        size_t i,
+        size_t vector_id,
+        size_t sq);
+
+/** set a single element "code" into a packed codes table
+ *
+ * @param vector_id       vector id
+ * @param sq       subquantizer (< nsq)
+ */
+void pq4_set_packed_element(
+        uint8_t* data,
+        uint8_t code,
+        size_t bbs,
+        size_t nsq,
+        size_t vector_id,
         size_t sq);
 
 /** Pack Look-up table for consumption by the kernel.
@@ -88,8 +101,9 @@ void pq4_pack_LUT(int nq, int nsq, const uint8_t* src, uint8_t* dest);
  * @param nsq     number of sub-quantizers (muliple of 2)
  * @param codes   packed codes array
  * @param LUT     packed look-up table
+ * @param scaler  scaler to scale the encoded norm
  */
-template <class ResultHandler>
+template <class ResultHandler, class Scaler>
 void pq4_accumulate_loop(
         int nq,
         size_t nb,
@@ -97,7 +111,8 @@ void pq4_accumulate_loop(
         int nsq,
         const uint8_t* codes,
         const uint8_t* LUT,
-        ResultHandler& res);
+        ResultHandler& res,
+        const Scaler& scaler);
 
 /* qbs versions, supported only for bbs=32.
  *
@@ -141,20 +156,22 @@ int pq4_pack_LUT_qbs_q_map(
 
 /** Run accumulation loop.
  *
- * @param qbs     4-bit encded number of queries
+ * @param qbs     4-bit encoded number of queries
  * @param nb      number of database codes (mutliple of bbs)
  * @param nsq     number of sub-quantizers
  * @param codes   encoded database vectors (packed)
  * @param LUT     look-up table (packed)
  * @param res     call-back for the resutls
+ * @param scaler  scaler to scale the encoded norm
  */
-template <class ResultHandler>
+template <class ResultHandler, class Scaler>
 void pq4_accumulate_loop_qbs(
         int qbs,
         size_t nb,
         int nsq,
         const uint8_t* codes,
         const uint8_t* LUT,
-        ResultHandler& res);
+        ResultHandler& res,
+        const Scaler& scaler);
 
 } // namespace faiss

data/vendor/faiss/faiss/impl/pq4_fast_scan_search_1.cpp

@@ -8,6 +8,7 @@
 #include <faiss/impl/pq4_fast_scan.h>
 
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/LookupTableScaler.h>
 #include <faiss/impl/simd_result_handlers.h>
 
 namespace faiss {
@@ -26,12 +27,13 @@ namespace {
  * writes results in a ResultHandler
  */
 
-template <int NQ, int BB, class ResultHandler>
+template <int NQ, int BB, class ResultHandler, class Scaler>
 void kernel_accumulate_block(
         int nsq,
         const uint8_t* codes,
         const uint8_t* LUT,
-        ResultHandler& res) {
+        ResultHandler& res,
+        const Scaler& scaler) {
     // distance accumulators
     simd16uint16 accu[NQ][BB][4];
 
@@ -44,7 +46,7 @@ void kernel_accumulate_block(
         }
     }
 
-    for (int sq = 0; sq < nsq; sq += 2) {
+    for (int sq = 0; sq < nsq - scaler.nscale; sq += 2) {
         simd32uint8 lut_cache[NQ];
         for (int q = 0; q < NQ; q++) {
             lut_cache[q] = simd32uint8(LUT);
@@ -72,6 +74,35 @@ void kernel_accumulate_block(
         }
     }
 
+    for (int sq = 0; sq < scaler.nscale; sq += 2) {
+        simd32uint8 lut_cache[NQ];
+        for (int q = 0; q < NQ; q++) {
+            lut_cache[q] = simd32uint8(LUT);
+            LUT += 32;
+        }
+
+        for (int b = 0; b < BB; b++) {
+            simd32uint8 c = simd32uint8(codes);
+            codes += 32;
+            simd32uint8 mask(15);
+            simd32uint8 chi = simd32uint8(simd16uint16(c) >> 4) & mask;
+            simd32uint8 clo = c & mask;
+
+            for (int q = 0; q < NQ; q++) {
+                simd32uint8 lut = lut_cache[q];
+
+                simd32uint8 res0 = scaler.lookup(lut, clo);
+                accu[q][b][0] += scaler.scale_lo(res0); // handle vectors 0..7
+                accu[q][b][1] += scaler.scale_hi(res0); // handle vectors 8..15
+
+                simd32uint8 res1 = scaler.lookup(lut, chi);
+                accu[q][b][2] += scaler.scale_lo(res1); // handle vectors 16..23
+                accu[q][b][3] +=
+                        scaler.scale_hi(res1); //  handle vectors 24..31
+            }
+        }
+    }
+
     for (int q = 0; q < NQ; q++) {
         for (int b = 0; b < BB; b++) {
             accu[q][b][0] -= accu[q][b][1] << 8;
@@ -85,17 +116,18 @@ void kernel_accumulate_block(
     }
 }
 
-template <int NQ, int BB, class ResultHandler>
+template <int NQ, int BB, class ResultHandler, class Scaler>
 void accumulate_fixed_blocks(
         size_t nb,
         int nsq,
         const uint8_t* codes,
         const uint8_t* LUT,
-        ResultHandler& res) {
+        ResultHandler& res,
+        const Scaler& scaler) {
     constexpr int bbs = 32 * BB;
     for (int64_t j0 = 0; j0 < nb; j0 += bbs) {
         FixedStorageHandler<NQ, 2 * BB> res2;
-        kernel_accumulate_block<NQ, BB>(nsq, codes, LUT, res2);
+        kernel_accumulate_block<NQ, BB>(nsq, codes, LUT, res2, scaler);
         res.set_block_origin(0, j0);
         res2.to_other_handler(res);
         codes += bbs * nsq / 2;
@@ -104,7 +136,7 @@ void accumulate_fixed_blocks(
 
 } // anonymous namespace
 
-template <class ResultHandler>
+template <class ResultHandler, class Scaler>
 void pq4_accumulate_loop(
         int nq,
         size_t nb,
@@ -112,15 +144,16 @@ void pq4_accumulate_loop(
         int nsq,
         const uint8_t* codes,
         const uint8_t* LUT,
-        ResultHandler& res) {
+        ResultHandler& res,
+        const Scaler& scaler) {
     FAISS_THROW_IF_NOT(is_aligned_pointer(codes));
     FAISS_THROW_IF_NOT(is_aligned_pointer(LUT));
     FAISS_THROW_IF_NOT(bbs % 32 == 0);
     FAISS_THROW_IF_NOT(nb % bbs == 0);
 
-#define DISPATCH(NQ, BB)                                           \
-    case NQ * 1000 + BB:                                           \
-        accumulate_fixed_blocks<NQ, BB>(nb, nsq, codes, LUT, res); \
+#define DISPATCH(NQ, BB)                                                   \
+    case NQ * 1000 + BB:                                                   \
+        accumulate_fixed_blocks<NQ, BB>(nb, nsq, codes, LUT, res, scaler); \
         break
 
     switch (nq * 1000 + bbs / 32) {
@@ -141,20 +174,28 @@ void pq4_accumulate_loop(
 
 // explicit template instantiations
 
-#define INSTANTIATE_ACCUMULATE(TH, C, with_id_map)         \
-    template void pq4_accumulate_loop<TH<C, with_id_map>>( \
-            int,                                           \
-            size_t,                                        \
-            int,                                           \
-            int,                                           \
-            const uint8_t*,                                \
-            const uint8_t*,                                \
-            TH<C, with_id_map>&);
-
-#define INSTANTIATE_3(C, with_id_map)                           \
-    INSTANTIATE_ACCUMULATE(SingleResultHandler, C, with_id_map) \
-    INSTANTIATE_ACCUMULATE(HeapHandler, C, with_id_map)         \
-    INSTANTIATE_ACCUMULATE(ReservoirHandler, C, with_id_map)
+#define INSTANTIATE_ACCUMULATE(TH, C, with_id_map, S)         \
+    template void pq4_accumulate_loop<TH<C, with_id_map>, S>( \
+            int,                                              \
+            size_t,                                           \
+            int,                                              \
+            int,                                              \
+            const uint8_t*,                                   \
+            const uint8_t*,                                   \
+            TH<C, with_id_map>&,                              \
+            const S&);
+
+using DS = DummyScaler;
+using NS = NormTableScaler;
+
+#define INSTANTIATE_3(C, with_id_map)                               \
+    INSTANTIATE_ACCUMULATE(SingleResultHandler, C, with_id_map, DS) \
+    INSTANTIATE_ACCUMULATE(HeapHandler, C, with_id_map, DS)         \
+    INSTANTIATE_ACCUMULATE(ReservoirHandler, C, with_id_map, DS)    \
+                                                                    \
+    INSTANTIATE_ACCUMULATE(SingleResultHandler, C, with_id_map, NS) \
+    INSTANTIATE_ACCUMULATE(HeapHandler, C, with_id_map, NS)         \
+    INSTANTIATE_ACCUMULATE(ReservoirHandler, C, with_id_map, NS)
 
 using Csi = CMax<uint16_t, int>;
 INSTANTIATE_3(Csi, false);