faiss 0.2.6 → 0.2.7

data/vendor/faiss/faiss/utils/partitioning.cpp

@@ -817,7 +817,7 @@ template uint16_t partition_fuzzy<CMax<uint16_t, int>>(
  * Histogram subroutines
  ******************************************************************/
 
-#ifdef __AVX2__
+#if defined(__AVX2__) || defined(__aarch64__)
 /// FIXME when MSB of uint16 is set
 // this code does not compile properly with GCC 7.4.0
 
@@ -833,7 +833,7 @@ simd32uint8 accu4to8(simd16uint16 a4) {
     simd16uint16 a8_0 = a4 & mask4;
     simd16uint16 a8_1 = (a4 >> 4) & mask4;
 
-    return simd32uint8(_mm256_hadd_epi16(a8_0.i, a8_1.i));
+    return simd32uint8(hadd(a8_0, a8_1));
 }
 
 simd16uint16 accu8to16(simd32uint8 a8) {
@@ -842,10 +842,10 @@ simd16uint16 accu8to16(simd32uint8 a8) {
     simd16uint16 a8_0 = simd16uint16(a8) & mask8;
     simd16uint16 a8_1 = (simd16uint16(a8) >> 8) & mask8;
 
-    return simd16uint16(_mm256_hadd_epi16(a8_0.i, a8_1.i));
+    return hadd(a8_0, a8_1);
 }
 
-static const simd32uint8 shifts(_mm256_setr_epi8(
+static const simd32uint8 shifts = simd32uint8::create<
         1,
         16,
         0,
@@ -877,7 +877,7 @@ static const simd32uint8 shifts(_mm256_setr_epi8(
         0,
         0,
         4,
-        64));
+        64>();
 
 // 2-bit accumulator: we can add only up to 3 elements
 // on output we return 2*4-bit results
@@ -937,7 +937,7 @@ simd16uint16 histogram_8(const uint16_t* data, Preproc pp, size_t n_in) {
     simd16uint16 a16lo = accu8to16(a8lo);
     simd16uint16 a16hi = accu8to16(a8hi);
 
-    simd16uint16 a16 = simd16uint16(_mm256_hadd_epi16(a16lo.i, a16hi.i));
+    simd16uint16 a16 = hadd(a16lo, a16hi);
 
     // the 2 lanes must still be combined
     return a16;
@@ -947,7 +947,7 @@ simd16uint16 histogram_8(const uint16_t* data, Preproc pp, size_t n_in) {
  * 16 bins
  ************************************************************/
 
-static const simd32uint8 shifts2(_mm256_setr_epi8(
+static const simd32uint8 shifts2 = simd32uint8::create<
         1,
         2,
         4,
@@ -955,7 +955,7 @@ static const simd32uint8 shifts2(_mm256_setr_epi8(
         16,
         32,
         64,
-        (char)128,
+        128,
         1,
         2,
         4,
@@ -963,7 +963,7 @@ static const simd32uint8 shifts2(_mm256_setr_epi8(
         16,
         32,
         64,
-        (char)128,
+        128,
         1,
         2,
         4,
@@ -971,7 +971,7 @@ static const simd32uint8 shifts2(_mm256_setr_epi8(
         16,
         32,
         64,
-        (char)128,
+        128,
         1,
         2,
         4,
@@ -979,19 +979,12 @@ static const simd32uint8 shifts2(_mm256_setr_epi8(
         16,
         32,
         64,
-        (char)128));
+        128>();
 
 simd32uint8 shiftr_16(simd32uint8 x, int n) {
     return simd32uint8(simd16uint16(x) >> n);
 }
 
-inline simd32uint8 combine_2x2(simd32uint8 a, simd32uint8 b) {
-    __m256i a1b0 = _mm256_permute2f128_si256(a.i, b.i, 0x21);
-    __m256i a0b1 = _mm256_blend_epi32(a.i, b.i, 0xF0);
-
-    return simd32uint8(a1b0) + simd32uint8(a0b1);
-}
-
 // 2-bit accumulator: we can add only up to 3 elements
 // on output we return 2*4-bit results
 template <int N, class Preproc>
@@ -1018,7 +1011,7 @@ void compute_accu2_16(
         // contains 0s for out-of-bounds elements
 
         simd16uint16 lt8 = (v >> 3) == simd16uint16(0);
-        lt8.i = _mm256_xor_si256(lt8.i, _mm256_set1_epi16(0xff00));
+        lt8 = lt8 ^ simd16uint16(0xff00);
 
         a1 = a1 & lt8;
 
@@ -1036,11 +1029,15 @@ void compute_accu2_16(
 simd32uint8 accu4to8_2(simd32uint8 a4_0, simd32uint8 a4_1) {
     simd32uint8 mask4(0x0f);
 
-    simd32uint8 a8_0 = combine_2x2(a4_0 & mask4, shiftr_16(a4_0, 4) & mask4);
+    simd16uint16 a8_0 = combine2x2(
+            (simd16uint16)(a4_0 & mask4),
+            (simd16uint16)(shiftr_16(a4_0, 4) & mask4));
 
-    simd32uint8 a8_1 = combine_2x2(a4_1 & mask4, shiftr_16(a4_1, 4) & mask4);
+    simd16uint16 a8_1 = combine2x2(
+            (simd16uint16)(a4_1 & mask4),
+            (simd16uint16)(shiftr_16(a4_1, 4) & mask4));
 
-    return simd32uint8(_mm256_hadd_epi16(a8_0.i, a8_1.i));
+    return simd32uint8(hadd(a8_0, a8_1));
 }
 
 template <class Preproc>
@@ -1079,10 +1076,9 @@ simd16uint16 histogram_16(const uint16_t* data, Preproc pp, size_t n_in) {
     simd16uint16 a16lo = accu8to16(a8lo);
     simd16uint16 a16hi = accu8to16(a8hi);
 
-    simd16uint16 a16 = simd16uint16(_mm256_hadd_epi16(a16lo.i, a16hi.i));
+    simd16uint16 a16 = hadd(a16lo, a16hi);
 
-    __m256i perm32 = _mm256_setr_epi32(0, 2, 4, 6, 1, 3, 5, 7);
-    a16.i = _mm256_permutevar8x32_epi32(a16.i, perm32);
+    a16 = simd16uint16{simd8uint32{a16}.unzip()};
 
     return a16;
 }

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

@@ -70,6 +70,13 @@ struct simd256bit {
         bin(bits);
         return std::string(bits);
     }
+
+    // Checks whether the other holds exactly the same bytes.
+    bool is_same_as(simd256bit other) const {
+        const __m256i pcmp = _mm256_cmpeq_epi32(i, other.i);
+        unsigned bitmask = _mm256_movemask_epi8(pcmp);
+        return (bitmask == 0xffffffffU);
+    }
 };
 
 /// vector of 16 elements in uint16
@@ -86,6 +93,41 @@ struct simd16uint16 : simd256bit {
 
     explicit simd16uint16(const uint16_t* x) : simd256bit((const void*)x) {}
 
+    explicit simd16uint16(
+            uint16_t u0,
+            uint16_t u1,
+            uint16_t u2,
+            uint16_t u3,
+            uint16_t u4,
+            uint16_t u5,
+            uint16_t u6,
+            uint16_t u7,
+            uint16_t u8,
+            uint16_t u9,
+            uint16_t u10,
+            uint16_t u11,
+            uint16_t u12,
+            uint16_t u13,
+            uint16_t u14,
+            uint16_t u15)
+            : simd256bit(_mm256_setr_epi16(
+                      u0,
+                      u1,
+                      u2,
+                      u3,
+                      u4,
+                      u5,
+                      u6,
+                      u7,
+                      u8,
+                      u9,
+                      u10,
+                      u11,
+                      u12,
+                      u13,
+                      u14,
+                      u15)) {}
+
     std::string elements_to_string(const char* fmt) const {
         uint16_t bytes[16];
         storeu((void*)bytes);
@@ -151,9 +193,19 @@ struct simd16uint16 : simd256bit {
         return simd16uint16(_mm256_or_si256(i, other.i));
     }
 
+    simd16uint16 operator^(simd256bit other) const {
+        return simd16uint16(_mm256_xor_si256(i, other.i));
+    }
+
     // returns binary masks
-    simd16uint16 operator==(simd256bit other) const {
-        return simd16uint16(_mm256_cmpeq_epi16(i, other.i));
+    friend simd16uint16 operator==(const simd256bit lhs, const simd256bit rhs) {
+        return simd16uint16(_mm256_cmpeq_epi16(lhs.i, rhs.i));
+    }
+
+    bool is_same(simd16uint16 other) const {
+        const __m256i pcmp = _mm256_cmpeq_epi16(i, other.i);
+        unsigned bitmask = _mm256_movemask_epi8(pcmp);
+        return (bitmask == 0xffffffffU);
     }
 
     simd16uint16 operator~() const {
@@ -255,6 +307,45 @@ inline uint32_t cmp_le32(simd16uint16 d0, simd16uint16 d1, simd16uint16 thr) {
     return ge;
 }
 
+inline simd16uint16 hadd(const simd16uint16& a, const simd16uint16& b) {
+    return simd16uint16(_mm256_hadd_epi16(a.i, b.i));
+}
+
+// Vectorized version of the following code:
+//   for (size_t i = 0; i < n; i++) {
+//      bool flag = (candidateValues[i] < currentValues[i]);
+//      minValues[i] = flag ? candidateValues[i] : currentValues[i];
+//      minIndices[i] = flag ? candidateIndices[i] : currentIndices[i];
+//      maxValues[i] = !flag ? candidateValues[i] : currentValues[i];
+//      maxIndices[i] = !flag ? candidateIndices[i] : currentIndices[i];
+//   }
+// Max indices evaluation is inaccurate in case of equal values (the index of
+// the last equal value is saved instead of the first one), but this behavior
+// saves instructions.
+//
+// Works in i16 mode in order to save instructions. One may
+// switch from i16 to u16.
+inline void cmplt_min_max_fast(
+        const simd16uint16 candidateValues,
+        const simd16uint16 candidateIndices,
+        const simd16uint16 currentValues,
+        const simd16uint16 currentIndices,
+        simd16uint16& minValues,
+        simd16uint16& minIndices,
+        simd16uint16& maxValues,
+        simd16uint16& maxIndices) {
+    // there's no lt instruction, so we'll need to emulate one
+    __m256i comparison = _mm256_cmpgt_epi16(currentValues.i, candidateValues.i);
+    comparison = _mm256_andnot_si256(comparison, _mm256_set1_epi16(-1));
+
+    minValues.i = _mm256_min_epi16(candidateValues.i, currentValues.i);
+    minIndices.i = _mm256_blendv_epi8(
+            candidateIndices.i, currentIndices.i, comparison);
+    maxValues.i = _mm256_max_epi16(candidateValues.i, currentValues.i);
+    maxIndices.i = _mm256_blendv_epi8(
+            currentIndices.i, candidateIndices.i, comparison);
+}
+
 // vector of 32 unsigned 8-bit integers
 struct simd32uint8 : simd256bit {
     simd32uint8() {}
@@ -265,6 +356,75 @@ struct simd32uint8 : simd256bit {
 
     explicit simd32uint8(uint8_t x) : simd256bit(_mm256_set1_epi8(x)) {}
 
+    template <
+            uint8_t _0,
+            uint8_t _1,
+            uint8_t _2,
+            uint8_t _3,
+            uint8_t _4,
+            uint8_t _5,
+            uint8_t _6,
+            uint8_t _7,
+            uint8_t _8,
+            uint8_t _9,
+            uint8_t _10,
+            uint8_t _11,
+            uint8_t _12,
+            uint8_t _13,
+            uint8_t _14,
+            uint8_t _15,
+            uint8_t _16,
+            uint8_t _17,
+            uint8_t _18,
+            uint8_t _19,
+            uint8_t _20,
+            uint8_t _21,
+            uint8_t _22,
+            uint8_t _23,
+            uint8_t _24,
+            uint8_t _25,
+            uint8_t _26,
+            uint8_t _27,
+            uint8_t _28,
+            uint8_t _29,
+            uint8_t _30,
+            uint8_t _31>
+    static simd32uint8 create() {
+        return simd32uint8(_mm256_setr_epi8(
+                (char)_0,
+                (char)_1,
+                (char)_2,
+                (char)_3,
+                (char)_4,
+                (char)_5,
+                (char)_6,
+                (char)_7,
+                (char)_8,
+                (char)_9,
+                (char)_10,
+                (char)_11,
+                (char)_12,
+                (char)_13,
+                (char)_14,
+                (char)_15,
+                (char)_16,
+                (char)_17,
+                (char)_18,
+                (char)_19,
+                (char)_20,
+                (char)_21,
+                (char)_22,
+                (char)_23,
+                (char)_24,
+                (char)_25,
+                (char)_26,
+                (char)_27,
+                (char)_28,
+                (char)_29,
+                (char)_30,
+                (char)_31));
+    }
+
     explicit simd32uint8(simd256bit x) : simd256bit(x) {}
 
     explicit simd32uint8(const uint8_t* x) : simd256bit((const void*)x) {}
@@ -359,6 +519,40 @@ struct simd8uint32 : simd256bit {
 
     explicit simd8uint32(const uint8_t* x) : simd256bit((const void*)x) {}
 
+    explicit simd8uint32(
+            uint32_t u0,
+            uint32_t u1,
+            uint32_t u2,
+            uint32_t u3,
+            uint32_t u4,
+            uint32_t u5,
+            uint32_t u6,
+            uint32_t u7)
+            : simd256bit(_mm256_setr_epi32(u0, u1, u2, u3, u4, u5, u6, u7)) {}
+
+    simd8uint32 operator+(simd8uint32 other) const {
+        return simd8uint32(_mm256_add_epi32(i, other.i));
+    }
+
+    simd8uint32 operator-(simd8uint32 other) const {
+        return simd8uint32(_mm256_sub_epi32(i, other.i));
+    }
+
+    simd8uint32& operator+=(const simd8uint32& other) {
+        i = _mm256_add_epi32(i, other.i);
+        return *this;
+    }
+
+    bool operator==(simd8uint32 other) const {
+        const __m256i pcmp = _mm256_cmpeq_epi32(i, other.i);
+        unsigned bitmask = _mm256_movemask_epi8(pcmp);
+        return (bitmask == 0xffffffffU);
+    }
+
+    bool operator!=(simd8uint32 other) const {
+        return !(*this == other);
+    }
+
     std::string elements_to_string(const char* fmt) const {
         uint32_t bytes[8];
         storeu((void*)bytes);
@@ -383,8 +577,49 @@ struct simd8uint32 : simd256bit {
     void set1(uint32_t x) {
         i = _mm256_set1_epi32((int)x);
     }
+
+    simd8uint32 unzip() const {
+        return simd8uint32(_mm256_permutevar8x32_epi32(
+                i, _mm256_setr_epi32(0, 2, 4, 6, 1, 3, 5, 7)));
+    }
 };
 
+// Vectorized version of the following code:
+//   for (size_t i = 0; i < n; i++) {
+//      bool flag = (candidateValues[i] < currentValues[i]);
+//      minValues[i] = flag ? candidateValues[i] : currentValues[i];
+//      minIndices[i] = flag ? candidateIndices[i] : currentIndices[i];
+//      maxValues[i] = !flag ? candidateValues[i] : currentValues[i];
+//      maxIndices[i] = !flag ? candidateIndices[i] : currentIndices[i];
+//   }
+// Max indices evaluation is inaccurate in case of equal values (the index of
+// the last equal value is saved instead of the first one), but this behavior
+// saves instructions.
+inline void cmplt_min_max_fast(
+        const simd8uint32 candidateValues,
+        const simd8uint32 candidateIndices,
+        const simd8uint32 currentValues,
+        const simd8uint32 currentIndices,
+        simd8uint32& minValues,
+        simd8uint32& minIndices,
+        simd8uint32& maxValues,
+        simd8uint32& maxIndices) {
+    // there's no lt instruction, so we'll need to emulate one
+    __m256i comparison = _mm256_cmpgt_epi32(currentValues.i, candidateValues.i);
+    comparison = _mm256_andnot_si256(comparison, _mm256_set1_epi32(-1));
+
+    minValues.i = _mm256_min_epi32(candidateValues.i, currentValues.i);
+    minIndices.i = _mm256_castps_si256(_mm256_blendv_ps(
+            _mm256_castsi256_ps(candidateIndices.i),
+            _mm256_castsi256_ps(currentIndices.i),
+            _mm256_castsi256_ps(comparison)));
+    maxValues.i = _mm256_max_epi32(candidateValues.i, currentValues.i);
+    maxIndices.i = _mm256_castps_si256(_mm256_blendv_ps(
+            _mm256_castsi256_ps(currentIndices.i),
+            _mm256_castsi256_ps(candidateIndices.i),
+            _mm256_castsi256_ps(comparison)));
+}
+
 struct simd8float32 : simd256bit {
     simd8float32() {}
 
@@ -394,7 +629,18 @@ struct simd8float32 : simd256bit {
 
     explicit simd8float32(float x) : simd256bit(_mm256_set1_ps(x)) {}
 
-    explicit simd8float32(const float* x) : simd256bit(_mm256_load_ps(x)) {}
+    explicit simd8float32(const float* x) : simd256bit(_mm256_loadu_ps(x)) {}
+
+    explicit simd8float32(
+            float f0,
+            float f1,
+            float f2,
+            float f3,
+            float f4,
+            float f5,
+            float f6,
+            float f7)
+            : simd256bit(_mm256_setr_ps(f0, f1, f2, f3, f4, f5, f6, f7)) {}
 
     simd8float32 operator*(simd8float32 other) const {
         return simd8float32(_mm256_mul_ps(f, other.f));
@@ -408,6 +654,22 @@ struct simd8float32 : simd256bit {
         return simd8float32(_mm256_sub_ps(f, other.f));
     }
 
+    simd8float32& operator+=(const simd8float32& other) {
+        f = _mm256_add_ps(f, other.f);
+        return *this;
+    }
+
+    bool operator==(simd8float32 other) const {
+        const __m256i pcmp =
+                _mm256_castps_si256(_mm256_cmp_ps(f, other.f, _CMP_EQ_OQ));
+        unsigned bitmask = _mm256_movemask_epi8(pcmp);
+        return (bitmask == 0xffffffffU);
+    }
+
+    bool operator!=(simd8float32 other) const {
+        return !(*this == other);
+    }
+
     std::string tostring() const {
         float tab[8];
         storeu((void*)tab);
@@ -439,6 +701,85 @@ inline simd8float32 fmadd(simd8float32 a, simd8float32 b, simd8float32 c) {
     return simd8float32(_mm256_fmadd_ps(a.f, b.f, c.f));
 }
 
+// The following primitive is a vectorized version of the following code
+// snippet:
+//   float lowestValue = HUGE_VAL;
+//   uint lowestIndex = 0;
+//   for (size_t i = 0; i < n; i++) {
+//     if (values[i] < lowestValue) {
+//       lowestValue = values[i];
+//       lowestIndex = i;
+//     }
+//   }
+// Vectorized version can be implemented via two operations: cmp and blend
+// with something like this:
+//   lowestValues = [HUGE_VAL; 8];
+//   lowestIndices = {0, 1, 2, 3, 4, 5, 6, 7};
+//   for (size_t i = 0; i < n; i += 8) {
+//     auto comparison = cmp(values + i, lowestValues);
+//     lowestValues = blend(
+//         comparison,
+//         values + i,
+//         lowestValues);
+//     lowestIndices = blend(
+//         comparison,
+//         i + {0, 1, 2, 3, 4, 5, 6, 7},
+//         lowestIndices);
+//     lowestIndices += {8, 8, 8, 8, 8, 8, 8, 8};
+//   }
+// The problem is that blend primitive needs very different instruction
+// order for AVX and ARM.
+// So, let's introduce a combination of these two in order to avoid
+// confusion for ppl who write in low-level SIMD instructions. Additionally,
+// these two ops (cmp and blend) are very often used together.
+inline void cmplt_and_blend_inplace(
+        const simd8float32 candidateValues,
+        const simd8uint32 candidateIndices,
+        simd8float32& lowestValues,
+        simd8uint32& lowestIndices) {
+    const __m256 comparison =
+            _mm256_cmp_ps(lowestValues.f, candidateValues.f, _CMP_LE_OS);
+    lowestValues.f = _mm256_min_ps(candidateValues.f, lowestValues.f);
+    lowestIndices.i = _mm256_castps_si256(_mm256_blendv_ps(
+            _mm256_castsi256_ps(candidateIndices.i),
+            _mm256_castsi256_ps(lowestIndices.i),
+            comparison));
+}
+
+// Vectorized version of the following code:
+//   for (size_t i = 0; i < n; i++) {
+//      bool flag = (candidateValues[i] < currentValues[i]);
+//      minValues[i] = flag ? candidateValues[i] : currentValues[i];
+//      minIndices[i] = flag ? candidateIndices[i] : currentIndices[i];
+//      maxValues[i] = !flag ? candidateValues[i] : currentValues[i];
+//      maxIndices[i] = !flag ? candidateIndices[i] : currentIndices[i];
+//   }
+// Max indices evaluation is inaccurate in case of equal values (the index of
+// the last equal value is saved instead of the first one), but this behavior
+// saves instructions.
+inline void cmplt_min_max_fast(
+        const simd8float32 candidateValues,
+        const simd8uint32 candidateIndices,
+        const simd8float32 currentValues,
+        const simd8uint32 currentIndices,
+        simd8float32& minValues,
+        simd8uint32& minIndices,
+        simd8float32& maxValues,
+        simd8uint32& maxIndices) {
+    const __m256 comparison =
+            _mm256_cmp_ps(currentValues.f, candidateValues.f, _CMP_LE_OS);
+    minValues.f = _mm256_min_ps(candidateValues.f, currentValues.f);
+    minIndices.i = _mm256_castps_si256(_mm256_blendv_ps(
+            _mm256_castsi256_ps(candidateIndices.i),
+            _mm256_castsi256_ps(currentIndices.i),
+            comparison));
+    maxValues.f = _mm256_max_ps(candidateValues.f, currentValues.f);
+    maxIndices.i = _mm256_castps_si256(_mm256_blendv_ps(
+            _mm256_castsi256_ps(currentIndices.i),
+            _mm256_castsi256_ps(candidateIndices.i),
+            comparison));
+}
+
 namespace {
 
 // get even float32's of a and b, interleaved