faiss 0.2.4 → 0.2.5

data/vendor/faiss/faiss/utils/extra_distances.cpp

@@ -14,6 +14,7 @@
 #include <cmath>
 
 #include <faiss/impl/AuxIndexStructures.h>
+#include <faiss/impl/DistanceComputer.h>
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/utils/utils.h>
 
@@ -89,18 +90,18 @@ void knn_extra_metrics_template(
 }
 
 template <class VD>
-struct ExtraDistanceComputer : DistanceComputer {
+struct ExtraDistanceComputer : FlatCodesDistanceComputer {
     VD vd;
     Index::idx_t nb;
     const float* q;
     const float* b;
 
-    float operator()(idx_t i) override {
-        return vd(q, b + i * vd.d);
+    float symmetric_dis(idx_t i, idx_t j) final {
+        return vd(b + j * vd.d, b + i * vd.d);
     }
 
-    float symmetric_dis(idx_t i, idx_t j) override {
-        return vd(b + j * vd.d, b + i * vd.d);
+    float distance_to_code(const uint8_t* code) final {
+        return vd(q, (float*)code);
     }
 
     ExtraDistanceComputer(
@@ -108,7 +109,11 @@ struct ExtraDistanceComputer : DistanceComputer {
             const float* xb,
             size_t nb,
             const float* q = nullptr)
-            : vd(vd), nb(nb), q(q), b(xb) {}
+            : FlatCodesDistanceComputer((uint8_t*)xb, vd.d * sizeof(float)),
+              vd(vd),
+              nb(nb),
+              q(q),
+              b(xb) {}
 
     void set_query(const float* x) override {
         q = x;
@@ -188,7 +193,7 @@ void knn_extra_metrics(
     }
 }
 
-DistanceComputer* get_extra_distance_computer(
+FlatCodesDistanceComputer* get_extra_distance_computer(
         size_t d,
         MetricType mt,
         float metric_arg,

data/vendor/faiss/faiss/utils/extra_distances.h

@@ -18,6 +18,8 @@
 
 namespace faiss {
 
+struct FlatCodesDistanceComputer;
+
 void pairwise_extra_distances(
         int64_t d,
         int64_t nq,
@@ -43,7 +45,7 @@ void knn_extra_metrics(
 
 /** get a DistanceComputer that refers to this type of distance and
  *  indexes a flat array of size nb */
-DistanceComputer* get_extra_distance_computer(
+FlatCodesDistanceComputer* get_extra_distance_computer(
         size_t d,
         MetricType mt,
         float metric_arg,

data/vendor/faiss/faiss/utils/fp16-fp16c.h

@@ -0,0 +1,21 @@
+#pragma once
+
+#include <immintrin.h>
+#include <cstdint>
+
+namespace faiss {
+
+inline uint16_t encode_fp16(float x) {
+    __m128 xf = _mm_set1_ps(x);
+    __m128i xi =
+            _mm_cvtps_ph(xf, _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC);
+    return _mm_cvtsi128_si32(xi) & 0xffff;
+}
+
+inline float decode_fp16(uint16_t x) {
+    __m128i xi = _mm_set1_epi16(x);
+    __m128 xf = _mm_cvtph_ps(xi);
+    return _mm_cvtss_f32(xf);
+}
+
+} // namespace faiss

data/vendor/faiss/faiss/utils/fp16-inl.h

@@ -0,0 +1,101 @@
+#pragma once
+
+#include <algorithm>
+#include <cstdint>
+
+namespace faiss {
+
+// non-intrinsic FP16 <-> FP32 code adapted from
+// https://github.com/ispc/ispc/blob/master/stdlib.ispc
+
+namespace {
+
+inline float floatbits(uint32_t x) {
+    void* xptr = &x;
+    return *(float*)xptr;
+}
+
+inline uint32_t intbits(float f) {
+    void* fptr = &f;
+    return *(uint32_t*)fptr;
+}
+
+} // namespace
+
+inline uint16_t encode_fp16(float f) {
+    // via Fabian "ryg" Giesen.
+    // https://gist.github.com/2156668
+    uint32_t sign_mask = 0x80000000u;
+    int32_t o;
+
+    uint32_t fint = intbits(f);
+    uint32_t sign = fint & sign_mask;
+    fint ^= sign;
+
+    // NOTE all the integer compares in this function can be safely
+    // compiled into signed compares since all operands are below
+    // 0x80000000. Important if you want fast straight SSE2 code (since
+    // there's no unsigned PCMPGTD).
+
+    // Inf or NaN (all exponent bits set)
+    // NaN->qNaN and Inf->Inf
+    // unconditional assignment here, will override with right value for
+    // the regular case below.
+    uint32_t f32infty = 255u << 23;
+    o = (fint > f32infty) ? 0x7e00u : 0x7c00u;
+
+    // (De)normalized number or zero
+    // update fint unconditionally to save the blending; we don't need it
+    // anymore for the Inf/NaN case anyway.
+
+    const uint32_t round_mask = ~0xfffu;
+    const uint32_t magic = 15u << 23;
+
+    // Shift exponent down, denormalize if necessary.
+    // NOTE This represents half-float denormals using single
+    // precision denormals.  The main reason to do this is that
+    // there's no shift with per-lane variable shifts in SSE*, which
+    // we'd otherwise need. It has some funky side effects though:
+    // - This conversion will actually respect the FTZ (Flush To Zero)
+    //   flag in MXCSR - if it's set, no half-float denormals will be
+    //   generated. I'm honestly not sure whether this is good or
+    //   bad. It's definitely interesting.
+    // - If the underlying HW doesn't support denormals (not an issue
+    //   with Intel CPUs, but might be a problem on GPUs or PS3 SPUs),
+    //   you will always get flush-to-zero behavior. This is bad,
+    //   unless you're on a CPU where you don't care.
+    // - Denormals tend to be slow. FP32 denormals are rare in
+    //   practice outside of things like recursive filters in DSP -
+    //   not a typical half-float application. Whether FP16 denormals
+    //   are rare in practice, I don't know. Whatever slow path your
+    //   HW may or may not have for denormals, this may well hit it.
+    float fscale = floatbits(fint & round_mask) * floatbits(magic);
+    fscale = std::min(fscale, floatbits((31u << 23) - 0x1000u));
+    int32_t fint2 = intbits(fscale) - round_mask;
+
+    if (fint < f32infty)
+        o = fint2 >> 13; // Take the bits!
+
+    return (o | (sign >> 16));
+}
+
+inline float decode_fp16(uint16_t h) {
+    // https://gist.github.com/2144712
+    // Fabian "ryg" Giesen.
+
+    const uint32_t shifted_exp = 0x7c00u << 13; // exponent mask after shift
+
+    int32_t o = ((int32_t)(h & 0x7fffu)) << 13; // exponent/mantissa bits
+    int32_t exp = shifted_exp & o;              // just the exponent
+    o += (int32_t)(127 - 15) << 23;             // exponent adjust
+
+    int32_t infnan_val = o + ((int32_t)(128 - 16) << 23);
+    int32_t zerodenorm_val =
+            intbits(floatbits(o + (1u << 23)) - floatbits(113u << 23));
+    int32_t reg_val = (exp == 0) ? zerodenorm_val : o;
+
+    int32_t sign_bit = ((int32_t)(h & 0x8000u)) << 16;
+    return floatbits(((exp == shifted_exp) ? infnan_val : reg_val) | sign_bit);
+}
+
+} // namespace faiss

data/vendor/faiss/faiss/utils/fp16.h

@@ -0,0 +1,11 @@
+#pragma once
+
+#include <cstdint>
+
+#include <faiss/impl/platform_macros.h>
+
+#if defined(__F16C__)
+#include <faiss/utils/fp16-fp16c.h>
+#else
+#include <faiss/utils/fp16-inl.h>
+#endif

data/vendor/faiss/faiss/utils/hamming-inl.h

@@ -9,6 +9,60 @@ namespace faiss {
 
 extern const uint8_t hamdis_tab_ham_bytes[256];
 
+/* Elementary Hamming distance computation: unoptimized  */
+template <size_t nbits, typename T>
+inline T hamming(const uint8_t* bs1, const uint8_t* bs2) {
+    const size_t nbytes = nbits / 8;
+    size_t i;
+    T h = 0;
+    for (i = 0; i < nbytes; i++) {
+        h += (T)hamdis_tab_ham_bytes[bs1[i] ^ bs2[i]];
+    }
+    return h;
+}
+
+/* Hamming distances for multiples of 64 bits */
+template <size_t nbits>
+inline hamdis_t hamming(const uint64_t* bs1, const uint64_t* bs2) {
+    const size_t nwords = nbits / 64;
+    size_t i;
+    hamdis_t h = 0;
+    for (i = 0; i < nwords; i++) {
+        h += popcount64(bs1[i] ^ bs2[i]);
+    }
+    return h;
+}
+
+/* specialized (optimized) functions */
+template <>
+inline hamdis_t hamming<64>(const uint64_t* pa, const uint64_t* pb) {
+    return popcount64(pa[0] ^ pb[0]);
+}
+
+template <>
+inline hamdis_t hamming<128>(const uint64_t* pa, const uint64_t* pb) {
+    return popcount64(pa[0] ^ pb[0]) + popcount64(pa[1] ^ pb[1]);
+}
+
+template <>
+inline hamdis_t hamming<256>(const uint64_t* pa, const uint64_t* pb) {
+    return popcount64(pa[0] ^ pb[0]) + popcount64(pa[1] ^ pb[1]) +
+            popcount64(pa[2] ^ pb[2]) + popcount64(pa[3] ^ pb[3]);
+}
+
+/* Hamming distances for multiple of 64 bits */
+inline hamdis_t hamming(
+        const uint64_t* bs1,
+        const uint64_t* bs2,
+        size_t nwords) {
+    hamdis_t h = 0;
+    for (size_t i = 0; i < nwords; i++) {
+        h += popcount64(bs1[i] ^ bs2[i]);
+    }
+    return h;
+}
+
+// BitstringWriter and BitstringReader functions
 inline BitstringWriter::BitstringWriter(uint8_t* code, size_t code_size)
         : code(code), code_size(code_size), i(0) {
     memset(code, 0, code_size);

data/vendor/faiss/faiss/utils/hamming.cpp

@@ -56,54 +56,6 @@ const uint8_t hamdis_tab_ham_bytes[256] = {
         4, 5, 5, 6, 5, 6, 6, 7, 3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7,
         4, 5, 5, 6, 5, 6, 6, 7, 5, 6, 6, 7, 6, 7, 7, 8};
 
-/* Elementary Hamming distance computation: unoptimized  */
-template <size_t nbits, typename T>
-T hamming(const uint8_t* bs1, const uint8_t* bs2) {
-    const size_t nbytes = nbits / 8;
-    size_t i;
-    T h = 0;
-    for (i = 0; i < nbytes; i++)
-        h += (T)hamdis_tab_ham_bytes[bs1[i] ^ bs2[i]];
-    return h;
-}
-
-/* Hamming distances for multiples of 64 bits */
-template <size_t nbits>
-hamdis_t hamming(const uint64_t* bs1, const uint64_t* bs2) {
-    const size_t nwords = nbits / 64;
-    size_t i;
-    hamdis_t h = 0;
-    for (i = 0; i < nwords; i++)
-        h += popcount64(bs1[i] ^ bs2[i]);
-    return h;
-}
-
-/* specialized (optimized) functions */
-template <>
-hamdis_t hamming<64>(const uint64_t* pa, const uint64_t* pb) {
-    return popcount64(pa[0] ^ pb[0]);
-}
-
-template <>
-hamdis_t hamming<128>(const uint64_t* pa, const uint64_t* pb) {
-    return popcount64(pa[0] ^ pb[0]) + popcount64(pa[1] ^ pb[1]);
-}
-
-template <>
-hamdis_t hamming<256>(const uint64_t* pa, const uint64_t* pb) {
-    return popcount64(pa[0] ^ pb[0]) + popcount64(pa[1] ^ pb[1]) +
-            popcount64(pa[2] ^ pb[2]) + popcount64(pa[3] ^ pb[3]);
-}
-
-/* Hamming distances for multiple of 64 bits */
-hamdis_t hamming(const uint64_t* bs1, const uint64_t* bs2, size_t nwords) {
-    size_t i;
-    hamdis_t h = 0;
-    for (i = 0; i < nwords; i++)
-        h += popcount64(bs1[i] ^ bs2[i]);
-    return h;
-}
-
 template <size_t nbits>
 void hammings(
         const uint64_t* bs1,

data/vendor/faiss/faiss/utils/ordered_key_value.h

@@ -46,6 +46,11 @@ struct CMin {
     inline static bool cmp(T a, T b) {
         return a < b;
     }
+    // Similar to cmp(), but also breaks ties
+    // by comparing the second pair of arguments.
+    inline static bool cmp2(T a1, T b1, TI a2, TI b2) {
+        return (a1 < b1) || ((a1 == b1) && (a2 < b2));
+    }
     inline static T neutral() {
         return std::numeric_limits<T>::lowest();
     }
@@ -64,6 +69,11 @@ struct CMax {
     inline static bool cmp(T a, T b) {
         return a > b;
     }
+    // Similar to cmp(), but also breaks ties
+    // by comparing the second pair of arguments.
+    inline static bool cmp2(T a1, T b1, TI a2, TI b2) {
+        return (a1 > b1) || ((a1 == b1) && (a2 > b2));
+    }
     inline static T neutral() {
         return std::numeric_limits<T>::max();
     }

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

@@ -284,6 +284,68 @@ void quantize_LUT_and_bias(
         *b_out = b;
 }
 
+void aq_quantize_LUT_and_bias(
+        size_t nprobe,
+        size_t M,
+        size_t ksub,
+        const float* LUT,
+        const float* bias,
+        size_t M_norm,
+        int norm_scale,
+        uint8_t* LUTq,
+        size_t M2,
+        uint16_t* biasq,
+        float* a_out,
+        float* b_out) {
+    float a, b;
+    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 += (i >= M - M_norm ? span * norm_scale : 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);
+
+    *a_out = a;
+    *b_out = b;
+}
+
+float aq_estimate_norm_scale(
+        size_t M,
+        size_t ksub,
+        size_t M_norm,
+        const float* LUT) {
+    float max_span_LUT = -HUGE_VAL;
+    for (int i = 0; i < M - M_norm; i++) {
+        float min = tab_min(LUT + i * ksub, ksub);
+        float span = tab_max(LUT + i * ksub, ksub) - min;
+        max_span_LUT = std::max(max_span_LUT, span);
+    }
+
+    float max_span_LUT_norm = -HUGE_VAL;
+    for (int i = M - M_norm; i < M; i++) {
+        float min = tab_min(LUT + i * ksub, ksub);
+        float span = tab_max(LUT + i * ksub, ksub) - min;
+        max_span_LUT_norm = std::max(max_span_LUT_norm, span);
+    }
+
+    return max_span_LUT_norm / max_span_LUT;
+}
+
 } // namespace quantize_lut
 
 } // namespace faiss

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

@@ -77,6 +77,26 @@ void quantize_LUT_and_bias(
         float* a_out = nullptr,
         float* b_out = nullptr);
 
+void aq_quantize_LUT_and_bias(
+        size_t nprobe,
+        size_t M,
+        size_t ksub,
+        const float* LUT,
+        const float* bias,
+        size_t M_norm,
+        int norm_scale,
+        uint8_t* LUTq,
+        size_t M2,
+        uint16_t* biasq,
+        float* a_out,
+        float* b_out);
+
+float aq_estimate_norm_scale(
+        size_t M,
+        size_t ksub,
+        size_t M_norm,
+        const float* LUT);
+
 } // namespace quantize_lut
 
 } // namespace faiss

data/vendor/faiss/faiss/utils/random.cpp

@@ -9,6 +9,23 @@
 
 #include <faiss/utils/random.h>
 
+extern "C" {
+int sgemm_(
+        const char* transa,
+        const char* transb,
+        FINTEGER* m,
+        FINTEGER* n,
+        FINTEGER* k,
+        const float* alpha,
+        const float* a,
+        FINTEGER* lda,
+        const float* b,
+        FINTEGER* ldb,
+        float* beta,
+        float* c,
+        FINTEGER* ldc);
+}
+
 namespace faiss {
 
 /**************************************************
@@ -165,4 +182,40 @@ void byte_rand(uint8_t* x, size_t n, int64_t seed) {
     }
 }
 
+void rand_smooth_vectors(size_t n, size_t d, float* x, int64_t seed) {
+    size_t d1 = 10;
+    std::vector<float> x1(n * d1);
+    float_randn(x1.data(), x1.size(), seed);
+    std::vector<float> rot(d1 * d);
+    float_rand(rot.data(), rot.size(), seed + 1);
+
+    { //
+        FINTEGER di = d, d1i = d1, ni = n;
+        float one = 1.0, zero = 0.0;
+        sgemm_("Not transposed",
+               "Not transposed", // natural order
+               &di,
+               &ni,
+               &d1i,
+               &one,
+               rot.data(),
+               &di, // rotation matrix
+               x1.data(),
+               &d1i, // second term
+               &zero,
+               x,
+               &di);
+    }
+
+    std::vector<float> scales(d);
+    float_rand(scales.data(), d, seed + 2);
+
+#pragma omp parallel for if (n * d > 10000)
+    for (int64_t i = 0; i < n; i++) {
+        for (size_t j = 0; j < d; j++) {
+            x[i * d + j] = sinf(x[i * d + j] * (scales[j] * 4 + 0.1));
+        }
+    }
+}
+
 } // namespace faiss

data/vendor/faiss/faiss/utils/random.h

@@ -54,4 +54,9 @@ void int64_rand_max(int64_t* x, size_t n, uint64_t max, int64_t seed);
 /* random permutation */
 void rand_perm(int* perm, size_t n, int64_t seed);
 
+/* Random set of vectors with intrinsic dimensionality 10 that is harder to
+ * index than a subspace of dim 10 but easier than uniform data in dimension d
+ * */
+void rand_smooth_vectors(size_t n, size_t d, float* x, int64_t seed);
+
 } // namespace faiss

data/vendor/faiss/faiss/utils/simdlib_avx2.h

@@ -111,6 +111,10 @@ struct simd16uint16 : simd256bit {
         i = _mm256_set1_epi16((short)x);
     }
 
+    simd16uint16 operator*(const simd16uint16& other) const {
+        return simd16uint16(_mm256_mullo_epi16(i, other.i));
+    }
+
     // shift must be known at compile time
     simd16uint16 operator>>(const int shift) const {
         return simd16uint16(_mm256_srli_epi16(i, shift));

data/vendor/faiss/faiss/utils/simdlib_emulated.h

@@ -120,6 +120,11 @@ struct simd16uint16 : simd256bit {
         }
     }
 
+    simd16uint16 operator*(const simd16uint16& other) const {
+        return binary_func(
+                *this, other, [](uint16_t a, uint16_t b) { return a * b; });
+    }
+
     // shift must be known at compile time
     simd16uint16 operator>>(const int shift) const {
         return unary_func(*this, [shift](uint16_t a) { return a >> shift; });
@@ -433,7 +438,7 @@ struct simd8uint32 : simd256bit {
 
     explicit simd8uint32(const simd256bit& x) : simd256bit(x) {}
 
-    explicit simd8uint32(const uint8_t* x) : simd256bit((const void*)x) {}
+    explicit simd8uint32(const uint32_t* x) : simd256bit((const void*)x) {}
 
     std::string elements_to_string(const char* fmt) const {
         char res[1000], *ptr = res;

data/vendor/faiss/faiss/utils/simdlib_neon.h

@@ -260,6 +260,11 @@ struct simd16uint16 {
         detail::simdlib::set1(data, &vdupq_n_u16, x);
     }
 
+    simd16uint16 operator*(const simd16uint16& other) const {
+        return simd16uint16{
+                detail::simdlib::binary_func(data, other.data, &vmulq_u16)};
+    }
+
     // shift must be known at compile time
     simd16uint16 operator>>(const int shift) const {
         switch (shift) {
@@ -641,8 +646,8 @@ inline simd32uint8 blendv(
     const uint8x16x2_t msb_mask = {
             vtstq_u8(mask.data.val[0], msb), vtstq_u8(mask.data.val[1], msb)};
     const uint8x16x2_t selected = {
-            vbslq_u8(msb_mask.val[0], a.data.val[0], b.data.val[0]),
-            vbslq_u8(msb_mask.val[1], a.data.val[1], b.data.val[1])};
+            vbslq_u8(msb_mask.val[0], b.data.val[0], a.data.val[0]),
+            vbslq_u8(msb_mask.val[1], b.data.val[1], a.data.val[1])};
     return simd32uint8{selected};
 }