faiss 0.3.0 → 0.3.2

data/vendor/faiss/faiss/utils/simdlib_avx512.h

@@ -0,0 +1,296 @@
+/**
+ * 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 <string>
+
+#include <immintrin.h>
+
+#include <faiss/impl/platform_macros.h>
+
+#include <faiss/utils/simdlib_avx2.h>
+
+namespace faiss {
+
+/** Simple wrapper around the AVX 512-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, and to give more readable names to the AVX intrinsics. It does not
+ * pretend to be exhausitve, functions are added as needed.
+ */
+
+/// 512-bit representation without interpretation as a vector
+struct simd512bit {
+    union {
+        __m512i i;
+        __m512 f;
+    };
+
+    simd512bit() {}
+
+    explicit simd512bit(__m512i i) : i(i) {}
+
+    explicit simd512bit(__m512 f) : f(f) {}
+
+    explicit simd512bit(const void* x)
+            : i(_mm512_loadu_si512((__m512i const*)x)) {}
+
+    // sets up a lower half of the register while keeping upper one as zero
+    explicit simd512bit(simd256bit lo)
+            : simd512bit(_mm512_inserti32x8(
+                      _mm512_castsi256_si512(lo.i),
+                      _mm256_setzero_si256(),
+                      1)) {}
+
+    // constructs from lower and upper halves
+    explicit simd512bit(simd256bit lo, simd256bit hi)
+            : simd512bit(_mm512_inserti32x8(
+                      _mm512_castsi256_si512(lo.i),
+                      hi.i,
+                      1)) {}
+
+    void clear() {
+        i = _mm512_setzero_si512();
+    }
+
+    void storeu(void* ptr) const {
+        _mm512_storeu_si512((__m512i*)ptr, i);
+    }
+
+    void loadu(const void* ptr) {
+        i = _mm512_loadu_si512((__m512i*)ptr);
+    }
+
+    void store(void* ptr) const {
+        _mm512_storeu_si512((__m512i*)ptr, i);
+    }
+
+    void bin(char bits[513]) const {
+        char bytes[64];
+        storeu((void*)bytes);
+        for (int i = 0; i < 512; i++) {
+            bits[i] = '0' + ((bytes[i / 8] >> (i % 8)) & 1);
+        }
+        bits[512] = 0;
+    }
+
+    std::string bin() const {
+        char bits[257];
+        bin(bits);
+        return std::string(bits);
+    }
+};
+
+/// vector of 32 elements in uint16
+struct simd32uint16 : simd512bit {
+    simd32uint16() {}
+
+    explicit simd32uint16(__m512i i) : simd512bit(i) {}
+
+    explicit simd32uint16(int x) : simd512bit(_mm512_set1_epi16(x)) {}
+
+    explicit simd32uint16(uint16_t x) : simd512bit(_mm512_set1_epi16(x)) {}
+
+    explicit simd32uint16(simd512bit x) : simd512bit(x) {}
+
+    explicit simd32uint16(const uint16_t* x) : simd512bit((const void*)x) {}
+
+    // sets up a lower half of the register
+    explicit simd32uint16(simd256bit lo) : simd512bit(lo) {}
+
+    // constructs from lower and upper halves
+    explicit simd32uint16(simd256bit lo, simd256bit hi) : simd512bit(lo, hi) {}
+
+    std::string elements_to_string(const char* fmt) const {
+        uint16_t bytes[32];
+        storeu((void*)bytes);
+        char res[2000];
+        char* ptr = res;
+        for (int i = 0; i < 32; i++) {
+            ptr += sprintf(ptr, fmt, bytes[i]);
+        }
+        // strip last ,
+        ptr[-1] = 0;
+        return std::string(res);
+    }
+
+    std::string hex() const {
+        return elements_to_string("%02x,");
+    }
+
+    std::string dec() const {
+        return elements_to_string("%3d,");
+    }
+
+    void set1(uint16_t x) {
+        i = _mm512_set1_epi16((short)x);
+    }
+
+    simd32uint16 operator*(const simd32uint16& other) const {
+        return simd32uint16(_mm512_mullo_epi16(i, other.i));
+    }
+
+    // shift must be known at compile time
+    simd32uint16 operator>>(const int shift) const {
+        return simd32uint16(_mm512_srli_epi16(i, shift));
+    }
+
+    // shift must be known at compile time
+    simd32uint16 operator<<(const int shift) const {
+        return simd32uint16(_mm512_slli_epi16(i, shift));
+    }
+
+    simd32uint16 operator+=(simd32uint16 other) {
+        i = _mm512_add_epi16(i, other.i);
+        return *this;
+    }
+
+    simd32uint16 operator-=(simd32uint16 other) {
+        i = _mm512_sub_epi16(i, other.i);
+        return *this;
+    }
+
+    simd32uint16 operator+(simd32uint16 other) const {
+        return simd32uint16(_mm512_add_epi16(i, other.i));
+    }
+
+    simd32uint16 operator-(simd32uint16 other) const {
+        return simd32uint16(_mm512_sub_epi16(i, other.i));
+    }
+
+    simd32uint16 operator&(simd512bit other) const {
+        return simd32uint16(_mm512_and_si512(i, other.i));
+    }
+
+    simd32uint16 operator|(simd512bit other) const {
+        return simd32uint16(_mm512_or_si512(i, other.i));
+    }
+
+    simd32uint16 operator^(simd512bit other) const {
+        return simd32uint16(_mm512_xor_si512(i, other.i));
+    }
+
+    simd32uint16 operator~() const {
+        return simd32uint16(_mm512_xor_si512(i, _mm512_set1_epi32(-1)));
+    }
+
+    simd16uint16 low() const {
+        return simd16uint16(_mm512_castsi512_si256(i));
+    }
+
+    simd16uint16 high() const {
+        return simd16uint16(_mm512_extracti32x8_epi32(i, 1));
+    }
+
+    // for debugging only
+    uint16_t operator[](int i) const {
+        ALIGNED(64) uint16_t tab[32];
+        store(tab);
+        return tab[i];
+    }
+
+    void accu_min(simd32uint16 incoming) {
+        i = _mm512_min_epu16(i, incoming.i);
+    }
+
+    void accu_max(simd32uint16 incoming) {
+        i = _mm512_max_epu16(i, incoming.i);
+    }
+};
+
+// decompose in 128-lanes: a = (a0, a1, a2, a3), b = (b0, b1, b2, b3)
+// return (a0 + a1 + a2 + a3, b0 + b1 + b2 + b3)
+inline simd16uint16 combine4x2(simd32uint16 a, simd32uint16 b) {
+    return combine2x2(a.low(), b.low()) + combine2x2(a.high(), b.high());
+}
+
+// vector of 32 unsigned 8-bit integers
+struct simd64uint8 : simd512bit {
+    simd64uint8() {}
+
+    explicit simd64uint8(__m512i i) : simd512bit(i) {}
+
+    explicit simd64uint8(int x) : simd512bit(_mm512_set1_epi8(x)) {}
+
+    explicit simd64uint8(uint8_t x) : simd512bit(_mm512_set1_epi8(x)) {}
+
+    // sets up a lower half of the register
+    explicit simd64uint8(simd256bit lo) : simd512bit(lo) {}
+
+    // constructs from lower and upper halves
+    explicit simd64uint8(simd256bit lo, simd256bit hi) : simd512bit(lo, hi) {}
+
+    explicit simd64uint8(simd512bit x) : simd512bit(x) {}
+
+    explicit simd64uint8(const uint8_t* x) : simd512bit((const void*)x) {}
+
+    std::string elements_to_string(const char* fmt) const {
+        uint8_t bytes[64];
+        storeu((void*)bytes);
+        char res[2000];
+        char* ptr = res;
+        for (int i = 0; i < 64; i++) {
+            ptr += sprintf(ptr, fmt, bytes[i]);
+        }
+        // strip last ,
+        ptr[-1] = 0;
+        return std::string(res);
+    }
+
+    std::string hex() const {
+        return elements_to_string("%02x,");
+    }
+
+    std::string dec() const {
+        return elements_to_string("%3d,");
+    }
+
+    void set1(uint8_t x) {
+        i = _mm512_set1_epi8((char)x);
+    }
+
+    simd64uint8 operator&(simd512bit other) const {
+        return simd64uint8(_mm512_and_si512(i, other.i));
+    }
+
+    simd64uint8 operator+(simd64uint8 other) const {
+        return simd64uint8(_mm512_add_epi8(i, other.i));
+    }
+
+    simd64uint8 lookup_4_lanes(simd64uint8 idx) const {
+        return simd64uint8(_mm512_shuffle_epi8(i, idx.i));
+    }
+
+    // extract + 0-extend lane
+    // this operation is slow (3 cycles)
+    simd32uint16 lane0_as_uint16() const {
+        __m256i x = _mm512_extracti32x8_epi32(i, 0);
+        return simd32uint16(_mm512_cvtepu8_epi16(x));
+    }
+
+    simd32uint16 lane1_as_uint16() const {
+        __m256i x = _mm512_extracti32x8_epi32(i, 1);
+        return simd32uint16(_mm512_cvtepu8_epi16(x));
+    }
+
+    simd64uint8 operator+=(simd64uint8 other) {
+        i = _mm512_add_epi8(i, other.i);
+        return *this;
+    }
+
+    // for debugging only
+    uint8_t operator[](int i) const {
+        ALIGNED(64) uint8_t tab[64];
+        store(tab);
+        return tab[i];
+    }
+};
+
+} // namespace faiss

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

@@ -168,9 +168,12 @@ static inline std::string elements_to_string(const char* fmt, const S& simd) {
     simd.store(bytes);
     char res[1000], *ptr = res;
     for (size_t i = 0; i < N; ++i) {
-        ptr += sprintf(ptr, fmt, bytes[i]);
+        int bytesWritten =
+                snprintf(ptr, sizeof(res) - (ptr - res), fmt, bytes[i]);
+        ptr += bytesWritten;
     }
-    // strip last ,
+    // The format usually contains a ',' separator so this is to remove the last
+    // separator.
     ptr[-1] = 0;
     return std::string(res);
 }
@@ -559,15 +562,13 @@ struct simd16uint16 {
     }
 
     // Checks whether the other holds exactly the same bytes.
-    bool is_same_as(simd16uint16 other) const {
-        const bool equal0 =
-                (vminvq_u16(vceqq_u16(data.val[0], other.data.val[0])) ==
-                 0xffff);
-        const bool equal1 =
-                (vminvq_u16(vceqq_u16(data.val[1], other.data.val[1])) ==
-                 0xffff);
-
-        return equal0 && equal1;
+    template <typename T>
+    bool is_same_as(T other) const {
+        const auto o = detail::simdlib::reinterpret_u16(other.data);
+        const auto equals = detail::simdlib::binary_func(data, o)
+                                    .template call<&vceqq_u16>();
+        const auto equal = vandq_u16(equals.val[0], equals.val[1]);
+        return vminvq_u16(equal) == 0xffffu;
     }
 
     simd16uint16 operator~() const {
@@ -689,13 +690,12 @@ inline void cmplt_min_max_fast(
         simd16uint16& minIndices,
         simd16uint16& maxValues,
         simd16uint16& maxIndices) {
-    const uint16x8x2_t comparison = uint16x8x2_t{
-            vcltq_u16(candidateValues.data.val[0], currentValues.data.val[0]),
-            vcltq_u16(candidateValues.data.val[1], currentValues.data.val[1])};
+    const uint16x8x2_t comparison =
+            detail::simdlib::binary_func(
+                    candidateValues.data, currentValues.data)
+                    .call<&vcltq_u16>();
 
-    minValues.data = uint16x8x2_t{
-            vminq_u16(candidateValues.data.val[0], currentValues.data.val[0]),
-            vminq_u16(candidateValues.data.val[1], currentValues.data.val[1])};
+    minValues = min(candidateValues, currentValues);
     minIndices.data = uint16x8x2_t{
             vbslq_u16(
                     comparison.val[0],
@@ -706,9 +706,7 @@ inline void cmplt_min_max_fast(
                     candidateIndices.data.val[1],
                     currentIndices.data.val[1])};
 
-    maxValues.data = uint16x8x2_t{
-            vmaxq_u16(candidateValues.data.val[0], currentValues.data.val[0]),
-            vmaxq_u16(candidateValues.data.val[1], currentValues.data.val[1])};
+    maxValues = max(candidateValues, currentValues);
     maxIndices.data = uint16x8x2_t{
             vbslq_u16(
                     comparison.val[0],
@@ -869,13 +867,13 @@ struct simd32uint8 {
     }
 
     // Checks whether the other holds exactly the same bytes.
-    bool is_same_as(simd32uint8 other) const {
-        const bool equal0 =
-                (vminvq_u8(vceqq_u8(data.val[0], other.data.val[0])) == 0xff);
-        const bool equal1 =
-                (vminvq_u8(vceqq_u8(data.val[1], other.data.val[1])) == 0xff);
-
-        return equal0 && equal1;
+    template <typename T>
+    bool is_same_as(T other) const {
+        const auto o = detail::simdlib::reinterpret_u8(other.data);
+        const auto equals = detail::simdlib::binary_func(data, o)
+                                    .template call<&vceqq_u8>();
+        const auto equal = vandq_u8(equals.val[0], equals.val[1]);
+        return vminvq_u8(equal) == 0xffu;
     }
 };
 
@@ -960,27 +958,28 @@ struct simd8uint32 {
         return *this;
     }
 
-    bool operator==(simd8uint32 other) const {
-        const auto equals = detail::simdlib::binary_func(data, other.data)
-                                    .call<&vceqq_u32>();
-        const auto equal = vandq_u32(equals.val[0], equals.val[1]);
-        return vminvq_u32(equal) == 0xffffffff;
+    simd8uint32 operator==(simd8uint32 other) const {
+        return simd8uint32{detail::simdlib::binary_func(data, other.data)
+                                   .call<&vceqq_u32>()};
     }
 
-    bool operator!=(simd8uint32 other) const {
-        return !(*this == other);
+    simd8uint32 operator~() const {
+        return simd8uint32{
+                detail::simdlib::unary_func(data).call<&vmvnq_u32>()};
     }
 
-    // Checks whether the other holds exactly the same bytes.
-    bool is_same_as(simd8uint32 other) const {
-        const bool equal0 =
-                (vminvq_u32(vceqq_u32(data.val[0], other.data.val[0])) ==
-                 0xffffffff);
-        const bool equal1 =
-                (vminvq_u32(vceqq_u32(data.val[1], other.data.val[1])) ==
-                 0xffffffff);
+    simd8uint32 operator!=(simd8uint32 other) const {
+        return ~(*this == other);
+    }
 
-        return equal0 && equal1;
+    // Checks whether the other holds exactly the same bytes.
+    template <typename T>
+    bool is_same_as(T other) const {
+        const auto o = detail::simdlib::reinterpret_u32(other.data);
+        const auto equals = detail::simdlib::binary_func(data, o)
+                                    .template call<&vceqq_u32>();
+        const auto equal = vandq_u32(equals.val[0], equals.val[1]);
+        return vminvq_u32(equal) == 0xffffffffu;
     }
 
     void clear() {
@@ -1053,13 +1052,14 @@ inline void cmplt_min_max_fast(
         simd8uint32& minIndices,
         simd8uint32& maxValues,
         simd8uint32& maxIndices) {
-    const uint32x4x2_t comparison = uint32x4x2_t{
-            vcltq_u32(candidateValues.data.val[0], currentValues.data.val[0]),
-            vcltq_u32(candidateValues.data.val[1], currentValues.data.val[1])};
-
-    minValues.data = uint32x4x2_t{
-            vminq_u32(candidateValues.data.val[0], currentValues.data.val[0]),
-            vminq_u32(candidateValues.data.val[1], currentValues.data.val[1])};
+    const uint32x4x2_t comparison =
+            detail::simdlib::binary_func(
+                    candidateValues.data, currentValues.data)
+                    .call<&vcltq_u32>();
+
+    minValues.data = detail::simdlib::binary_func(
+                             candidateValues.data, currentValues.data)
+                             .call<&vminq_u32>();
     minIndices.data = uint32x4x2_t{
             vbslq_u32(
                     comparison.val[0],
@@ -1070,9 +1070,9 @@ inline void cmplt_min_max_fast(
                     candidateIndices.data.val[1],
                     currentIndices.data.val[1])};
 
-    maxValues.data = uint32x4x2_t{
-            vmaxq_u32(candidateValues.data.val[0], currentValues.data.val[0]),
-            vmaxq_u32(candidateValues.data.val[1], currentValues.data.val[1])};
+    maxValues.data = detail::simdlib::binary_func(
+                             candidateValues.data, currentValues.data)
+                             .call<&vmaxq_u32>();
     maxIndices.data = uint32x4x2_t{
             vbslq_u32(
                     comparison.val[0],
@@ -1167,28 +1167,25 @@ struct simd8float32 {
         return *this;
     }
 
-    bool operator==(simd8float32 other) const {
-        const auto equals =
+    simd8uint32 operator==(simd8float32 other) const {
+        return simd8uint32{
                 detail::simdlib::binary_func<::uint32x4x2_t>(data, other.data)
-                        .call<&vceqq_f32>();
-        const auto equal = vandq_u32(equals.val[0], equals.val[1]);
-        return vminvq_u32(equal) == 0xffffffff;
+                        .call<&vceqq_f32>()};
     }
 
-    bool operator!=(simd8float32 other) const {
-        return !(*this == other);
+    simd8uint32 operator!=(simd8float32 other) const {
+        return ~(*this == other);
     }
 
     // Checks whether the other holds exactly the same bytes.
-    bool is_same_as(simd8float32 other) const {
-        const bool equal0 =
-                (vminvq_u32(vceqq_f32(data.val[0], other.data.val[0])) ==
-                 0xffffffff);
-        const bool equal1 =
-                (vminvq_u32(vceqq_f32(data.val[1], other.data.val[1])) ==
-                 0xffffffff);
-
-        return equal0 && equal1;
+    template <typename T>
+    bool is_same_as(T other) const {
+        const auto o = detail::simdlib::reinterpret_f32(other.data);
+        const auto equals =
+                detail::simdlib::binary_func<::uint32x4x2_t>(data, o)
+                        .template call<&vceqq_f32>();
+        const auto equal = vandq_u32(equals.val[0], equals.val[1]);
+        return vminvq_u32(equal) == 0xffffffffu;
     }
 
     std::string tostring() const {
@@ -1302,13 +1299,14 @@ inline void cmplt_min_max_fast(
         simd8uint32& minIndices,
         simd8float32& maxValues,
         simd8uint32& maxIndices) {
-    const uint32x4x2_t comparison = uint32x4x2_t{
-            vcltq_f32(candidateValues.data.val[0], currentValues.data.val[0]),
-            vcltq_f32(candidateValues.data.val[1], currentValues.data.val[1])};
-
-    minValues.data = float32x4x2_t{
-            vminq_f32(candidateValues.data.val[0], currentValues.data.val[0]),
-            vminq_f32(candidateValues.data.val[1], currentValues.data.val[1])};
+    const uint32x4x2_t comparison =
+            detail::simdlib::binary_func<::uint32x4x2_t>(
+                    candidateValues.data, currentValues.data)
+                    .call<&vcltq_f32>();
+
+    minValues.data = detail::simdlib::binary_func(
+                             candidateValues.data, currentValues.data)
+                             .call<&vminq_f32>();
     minIndices.data = uint32x4x2_t{
             vbslq_u32(
                     comparison.val[0],
@@ -1319,9 +1317,9 @@ inline void cmplt_min_max_fast(
                     candidateIndices.data.val[1],
                     currentIndices.data.val[1])};
 
-    maxValues.data = float32x4x2_t{
-            vmaxq_f32(candidateValues.data.val[0], currentValues.data.val[0]),
-            vmaxq_f32(candidateValues.data.val[1], currentValues.data.val[1])};
+    maxValues.data = detail::simdlib::binary_func(
+                             candidateValues.data, currentValues.data)
+                             .call<&vmaxq_f32>();
     maxIndices.data = uint32x4x2_t{
             vbslq_u32(
                     comparison.val[0],