faiss 0.1.3 → 0.2.0

data/vendor/faiss/faiss/utils/partitioning.h

@@ -0,0 +1,69 @@
+/**
+ * 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 <stdint.h>
+#include <stdio.h>
+
+#include <faiss/impl/platform_macros.h>
+
+namespace faiss {
+
+
+/** partitions the table into 0:q and q:n where all elements above q are >= all
+ * elements below q (for C = CMax, for CMin comparisons are reversed)
+ *
+ * Returns the partition threshold. The elements q:n are destroyed on output.
+ */
+template<class C>
+typename C::T partition_fuzzy(
+    typename C::T *vals, typename C::TI * ids, size_t n,
+    size_t q_min, size_t q_max, size_t * q_out);
+
+/** simplified interface for when the parition is not fuzzy */
+template<class C>
+inline typename C::T partition(
+    typename C::T *vals, typename C::TI * ids, size_t n,
+    size_t q)
+{
+    return partition_fuzzy<C>(vals, ids, n, q, q, nullptr);
+}
+
+/** low level SIMD histogramming functions */
+
+/** 8-bin histogram of (x - min) >> shift
+ * values outside the range are ignored.
+ * the data table should be aligned on 32 bytes */
+void simd_histogram_8(
+    const uint16_t *data, int n,
+    uint16_t min, int shift,
+    int *hist);
+
+/** same for 16-bin histogram */
+void simd_histogram_16(
+    const uint16_t *data, int n,
+    uint16_t min, int shift,
+    int *hist);
+
+
+struct PartitionStats {
+    uint64_t bissect_cycles;
+    uint64_t compress_cycles;
+
+    PartitionStats () {reset (); }
+    void reset ();
+};
+
+// global var that collects them all
+FAISS_API extern PartitionStats partition_stats;
+
+
+
+} // namespace faiss
+

data/vendor/faiss/faiss/utils/quantize_lut.cpp

@@ -0,0 +1,277 @@
+/**
+ * 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 <faiss/utils/quantize_lut.h>
+
+#include <cmath>
+#include <cstring>
+#include <vector>
+#include <algorithm>
+
+#include <faiss/impl/FaissAssert.h>
+
+
+namespace faiss {
+
+
+namespace quantize_lut {
+
+
+/******************************************************
+ * Quantize look-up tables
+ ******************************************************/
+
+namespace {
+
+float round_uint8_and_mul(float *tab, size_t n) {
+    float max = 0;
+    for(int i = 0; i < n; i++) {
+        if(fabs(tab[i]) > max) {
+            max = fabs(tab[i]);
+        }
+    }
+    float multiplier = 127 / max;
+    for(int i = 0; i < n; i++) {
+        tab[i] = floorf(tab[i] * multiplier + 128);
+    }
+    return multiplier;
+}
+
+// there can be NaNs in tables, they should be ignored
+float tab_min(const float *tab, size_t n) {
+    float min = HUGE_VAL;
+    for(int i = 0; i < n; i++) {
+        if (tab[i] < min) min = tab[i];
+    }
+    return min;
+}
+
+float tab_max(const float *tab, size_t n) {
+    float max = -HUGE_VAL;
+    for(int i = 0; i < n; i++) {
+        if (tab[i] > max) max = tab[i];
+    }
+    return max;
+}
+
+void round_tab(float *tab, size_t n, float a, float bi) {
+    for(int i = 0; i < n; i++) {
+        tab[i] = floorf((tab[i] - bi) * a + 0.5);
+    }
+}
+
+template<typename T>
+void round_tab(const float *tab, size_t n, float a, float bi, T *tab_out) {
+    for(int i = 0; i < n; i++) {
+        tab_out[i] = (T)floorf((tab[i] - bi) * a + 0.5);
+    }
+}
+
+
+
+} // anonymous namespace
+
+void round_uint8_per_column(
+        float *tab, size_t n, size_t d,
+        float *a_out, float *b_out)
+{
+    float max_span = 0;
+    std::vector<float> mins(n);
+    for(int i = 0; i < n; i++) {
+        mins[i] = tab_min(tab + i * d, d);
+        float span = tab_max(tab + i * d, d) - mins[i];
+        if(span > max_span) {
+            max_span = span;
+        }
+    }
+    float a = 255 / max_span;
+    float b = 0;
+    for(int i = 0; i < n; i++) {
+        b += mins[i];
+        round_tab(tab + i * d, d, a, mins[i]);
+    }
+    if (a_out) *a_out = a;
+    if (b_out) *b_out = b;
+}
+
+void round_uint8_per_column_multi(
+        float *tab, size_t m, size_t n, size_t d,
+        float *a_out, float *b_out)
+{
+    float max_span = 0;
+    std::vector<float> mins(n);
+    for(int i = 0; i < n; i++) {
+        float min_i = HUGE_VAL;
+        float max_i = -HUGE_VAL;
+        for(int j = 0; j < m; j++) {
+            min_i = std::min(min_i, tab_min(tab + (j * n + i) * d, d));
+            max_i = std::max(max_i, tab_max(tab + (j * n + i) * d, d));
+        }
+        mins[i] = min_i;
+        float span = max_i - min_i;
+        if(span > max_span) {
+            max_span = span;
+        }
+    }
+    float a = 255 / max_span;
+    float b = 0;
+    for(int i = 0; i < n; i++) {
+        b += mins[i];
+        for(int j = 0; j < m; j++) {
+            round_tab(tab + (j * n + i) * d, d, a, mins[i]);
+        }
+    }
+    if (a_out) *a_out = a;
+    if (b_out) *b_out = b;
+}
+
+
+// translation of
+// https://github.com/fairinternal/faiss_improvements/blob/7122c3cc6ddb0a371d8aa6f1309cd8bcf2335e61/LUT_quantization.ipynb
+void quantize_LUT_and_bias(
+        size_t nprobe, size_t M, size_t ksub,
+        bool lut_is_3d,
+        const float *LUT,
+        const float *bias,
+        uint8_t *LUTq, size_t M2,
+        uint16_t *biasq,
+        float *a_out, float *b_out)
+{
+    float a, b;
+    if (!bias) {
+        FAISS_THROW_IF_NOT(!lut_is_3d);
+        std::vector<float> mins(M);
+        float max_span_LUT = -HUGE_VAL, max_span_dis = 0;
+        b = 0;
+        for(int i = 0; i < M; i++) {
+            mins[i] = tab_min(LUT + i * ksub, ksub);
+            float span = tab_max(LUT + i * ksub, ksub) - mins[i];
+            max_span_LUT = std::max(max_span_LUT, span);
+            max_span_dis += span;
+            b += mins[i];
+        }
+        a = std::min(255 / max_span_LUT, 65535 / max_span_dis);
+
+        for(int i = 0; i < M; i++) {
+            round_tab(LUT + i * ksub, ksub, a, mins[i], LUTq + i * ksub);
+        }
+        memset(LUTq + M * ksub, 0, ksub * (M2 - M));
+    } else if (!lut_is_3d) {
+        std::vector<float> mins(M);
+        float max_span_LUT = -HUGE_VAL, max_span_dis;
+        float bias_min = tab_min(bias, nprobe);
+        float bias_max = tab_max(bias, nprobe);
+        max_span_dis = bias_max - bias_min;
+        b = 0;
+        for(int i = 0; i < M; i++) {
+            mins[i] = tab_min(LUT + i * ksub, ksub);
+            float span = tab_max(LUT + i * ksub, ksub) - mins[i];
+            max_span_LUT = std::max(max_span_LUT, span);
+            max_span_dis += span;
+            b += mins[i];
+        }
+        a = std::min(255 / max_span_LUT, 65535 / max_span_dis);
+        b += bias_min;
+
+        for(int i = 0; i < M; i++) {
+            round_tab(LUT + i * ksub, ksub, a, mins[i], LUTq + i * ksub);
+        }
+        memset(LUTq + M * ksub, 0, ksub * (M2 - M));
+        round_tab(bias, nprobe, a, bias_min, biasq);
+
+    } else if (biasq) {
+        // LUT is 3D
+        std::vector<float> mins(nprobe * M);
+        std::vector<float> bias2(nprobe);
+        float bias_min = tab_min(bias, nprobe);
+        float max_span_LUT = -HUGE_VAL, max_span_dis = -HUGE_VAL;
+
+        b = HUGE_VAL;
+        size_t ij = 0;
+        for (int j = 0; j < nprobe; j++) {
+            float max_span_dis_j = bias[j] - bias_min;
+            float b2j = bias[j];
+            for(int i = 0; i < M; i++) {
+                mins[ij] = tab_min(LUT + ij * ksub, ksub);
+                float span = tab_max(LUT + ij * ksub, ksub) - mins[ij];
+                max_span_LUT = std::max(max_span_LUT, span);
+                max_span_dis_j += span;
+                b2j += mins[ij];
+                ij++;
+            }
+            max_span_dis = std::max(max_span_dis, max_span_dis_j);
+            bias2[j] = b2j;
+            b = std::min(b, b2j);
+        }
+
+        a = std::min(255 / max_span_LUT, 65535 / max_span_dis);
+
+        ij = 0;
+        size_t ij_2 = 0;
+        for (int j = 0; j < nprobe; j++) {
+            for(int i = 0; i < M; i++) {
+                round_tab(LUT + ij * ksub, ksub, a, mins[ij], LUTq + ij_2 * ksub);
+                ij++; ij_2++;
+            }
+            memset(LUTq + ij_2 * ksub, 0, ksub * (M2 - M));
+            ij_2 += M2 - M;
+        }
+
+        round_tab(bias2.data(), nprobe, a, b, biasq);
+
+    } else { // !biasq
+        // then we integrate the bias into the LUTs
+        std::vector<float> LUT2_storage(nprobe * M * ksub);
+        float *LUT2 = LUT2_storage.data();
+        size_t ijc = 0;
+        for (int j = 0; j < nprobe; j++) {
+            float bias_j = bias[j] / M;
+            for(int i = 0; i < M; i++) {
+                for (int c = 0; c < ksub; c++) {
+                    LUT2[ijc] = LUT[ijc] + bias_j;
+                    ijc++;
+                }
+            }
+        }
+        std::vector<float> mins(M, HUGE_VAL), maxs(M, -HUGE_VAL);
+        size_t ij = 0;
+        for (int j = 0; j < nprobe; j++) {
+            for(int i = 0; i < M; i++) {
+                mins[i] = std::min(mins[i], tab_min(LUT2 + ij * ksub, ksub));
+                maxs[i] = std::max(maxs[i], tab_max(LUT2 + ij * ksub, ksub));
+                ij++;
+            }
+        }
+
+        float max_span = -HUGE_VAL;
+        b = 0;
+        for(int i = 0; i < M; i++) {
+            float span = maxs[i] - mins[i];
+            max_span = std::max(max_span, span);
+            b += mins[i];
+        }
+        a = 255 / max_span;
+        ij = 0;
+        size_t ij_2 = 0;
+        for (int j = 0; j < nprobe; j++) {
+            for(int i = 0; i < M; i++) {
+                round_tab(LUT2 + ij * ksub, ksub, a, mins[i], LUTq + ij_2 * ksub);
+                ij++; ij_2++;
+            }
+            memset(LUTq + ij_2 * ksub, 0, ksub * (M2 - M));
+            ij_2 += M2 - M;
+        }
+    }
+    if (a_out) *a_out = a;
+    if (b_out) *b_out = b;
+}
+
+
+} // namespace quantize_lut
+
+} // namespace faiss

data/vendor/faiss/faiss/utils/quantize_lut.h

@@ -0,0 +1,80 @@
+/**
+ * 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 <cstdio>
+#include <cstdint>
+
+namespace faiss {
+
+/** Functions to quantize PQ floating-point Look Up Tables (LUT) to uint8, and
+ * biases to uint16. The accumulation is supposed to take place in uint16.
+ * The quantization coefficients are float (a, b) such that
+ *
+ *      original_value = quantized_value * a / b
+ *
+ * The hardest part of the quantization is with multiple LUTs that need to be
+ * added up together. In that case, coefficient a has to be chosen so that
+ * the sum fits in a uint16 accumulator.
+ */
+
+namespace quantize_lut {
+
+/* affine quantizer, a and b are the affine coefficients, marginalize over d
+ *
+ * @param tab input/output, size (n, d)
+ */
+void round_uint8_per_column(
+        float *tab, size_t n, size_t d,
+        float *a_out = nullptr,
+        float *b_out = nullptr
+);
+
+
+/* affine quantizer, a and b are the affine coefficients
+ *
+ * @param tab input/output, size (m, n, d)
+ */
+void round_uint8_per_column_multi(
+        float *tab, size_t m, size_t n, size_t d,
+        float *a_out = nullptr, float *b_out = nullptr);
+
+/** LUT quantization to uint8 and bias to uint16.
+ *
+ * (nprobe, M, ksub, lut_is_3d) determine the size of the the LUT
+ *
+ *  LUT input:
+ *  - 2D size (M, ksub): single matrix per probe (lut_is_3d=false)
+ *  - 3D size (nprobe, M, ksub): separate LUT per probe (lut_is_3d=true)
+ *  bias input:
+ *  - nullptr: bias is 0
+ *  - size (nprobe): one bias per probe
+ *  Output:
+ *  - LUTq uint8 version of the LUT (M size is rounded up to M2)
+ *  - biasq (or nullptr): uint16 version of the LUT
+ *  - a, b: scalars to approximate the true distance
+ */
+
+void quantize_LUT_and_bias(
+        size_t nprobe, size_t M, size_t ksub,
+        bool lut_is_3d,
+        const float *LUT,
+        const float *bias,
+        uint8_t *LUTq, size_t M2,
+        uint16_t *biasq,
+        float *a_out = nullptr, float *b_out = nullptr
+);
+
+
+} // namespace quantize_lut
+
+} // namespace faiss
+
+

data/vendor/faiss/faiss/utils/simdlib.h

@@ -0,0 +1,31 @@
+/**
+ * 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
+
+
+
+/** Abstractions for 256-bit registers
+ *
+ * The objective is to separate the different interpretations of the same
+ * registers (as a vector of uint8, uint16 or uint32), to provide printing
+ * functions.
+ */
+
+#ifdef __AVX2__
+
+#include <faiss/utils/simdlib_avx2.h>
+
+#else
+
+// emulated = all operations are implemented as scalars
+#include <faiss/utils/simdlib_emulated.h>
+
+// FIXME: make a SSE version
+// is this ever going to happen? We will probably rather implement AVX512
+
+#endif