faiss 0.3.0 → 0.3.2

data/vendor/faiss/faiss/impl/ScalarQuantizer.h

@@ -32,6 +32,9 @@ struct ScalarQuantizer : Quantizer {
         QT_fp16,
         QT_8bit_direct, ///< fast indexing of uint8s
         QT_6bit,        ///< 6 bits per component
+        QT_bf16,
+        QT_8bit_direct_signed, ///< fast indexing of signed int8s ranging from
+                               ///< [-128 to 127]
     };
 
     QuantizerType qtype = QT_8bit;
@@ -65,14 +68,6 @@ struct ScalarQuantizer : Quantizer {
 
     void train(size_t n, const float* x) override;
 
-    /// Used by an IVF index to train based on the residuals
-    void train_residual(
-            size_t n,
-            const float* x,
-            Index* quantizer,
-            bool by_residual,
-            bool verbose);
-
     /** Encode a set of vectors
      *
      * @param x      vectors to encode, size n * d

data/vendor/faiss/faiss/impl/code_distance/code_distance-avx2.h

@@ -13,25 +13,223 @@
 
 #include <type_traits>
 
+#include <faiss/impl/ProductQuantizer.h>
 #include <faiss/impl/code_distance/code_distance-generic.h>
 
+// https://gcc.gnu.org/bugzilla/show_bug.cgi?id=78782
+#if defined(__GNUC__) && __GNUC__ < 9
+#define _mm_loadu_si64(x) (_mm_loadl_epi64((__m128i_u*)x))
+#endif
+
 namespace {
 
+inline float horizontal_sum(const __m128 v) {
+    const __m128 v0 = _mm_shuffle_ps(v, v, _MM_SHUFFLE(0, 0, 3, 2));
+    const __m128 v1 = _mm_add_ps(v, v0);
+    __m128 v2 = _mm_shuffle_ps(v1, v1, _MM_SHUFFLE(0, 0, 0, 1));
+    const __m128 v3 = _mm_add_ps(v1, v2);
+    return _mm_cvtss_f32(v3);
+}
+
 // Computes a horizontal sum over an __m256 register
-inline float horizontal_sum(const __m256 reg) {
-    const __m256 h0 = _mm256_hadd_ps(reg, reg);
-    const __m256 h1 = _mm256_hadd_ps(h0, h0);
+inline float horizontal_sum(const __m256 v) {
+    const __m128 v0 =
+            _mm_add_ps(_mm256_castps256_ps128(v), _mm256_extractf128_ps(v, 1));
+    return horizontal_sum(v0);
+}
+
+// processes a single code for M=4, ksub=256, nbits=8
+float inline distance_single_code_avx2_pqdecoder8_m4(
+        // precomputed distances, layout (4, 256)
+        const float* sim_table,
+        const uint8_t* code) {
+    float result = 0;
+
+    const float* tab = sim_table;
+    constexpr size_t ksub = 1 << 8;
+
+    const __m128i vksub = _mm_set1_epi32(ksub);
+    __m128i offsets_0 = _mm_setr_epi32(0, 1, 2, 3);
+    offsets_0 = _mm_mullo_epi32(offsets_0, vksub);
+
+    // accumulators of partial sums
+    __m128 partialSum;
+
+    // load 4 uint8 values
+    const __m128i mm1 = _mm_cvtsi32_si128(*((const int32_t*)code));
+    {
+        // convert uint8 values (low part of __m128i) to int32
+        // values
+        const __m128i idx1 = _mm_cvtepu8_epi32(mm1);
+
+        // add offsets
+        const __m128i indices_to_read_from = _mm_add_epi32(idx1, offsets_0);
+
+        // gather 8 values, similar to 8 operations of tab[idx]
+        __m128 collected =
+                _mm_i32gather_ps(tab, indices_to_read_from, sizeof(float));
+
+        // collect partial sums
+        partialSum = collected;
+    }
+
+    // horizontal sum for partialSum
+    result = horizontal_sum(partialSum);
+    return result;
+}
+
+// processes a single code for M=8, ksub=256, nbits=8
+float inline distance_single_code_avx2_pqdecoder8_m8(
+        // precomputed distances, layout (8, 256)
+        const float* sim_table,
+        const uint8_t* code) {
+    float result = 0;
+
+    const float* tab = sim_table;
+    constexpr size_t ksub = 1 << 8;
+
+    const __m256i vksub = _mm256_set1_epi32(ksub);
+    __m256i offsets_0 = _mm256_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7);
+    offsets_0 = _mm256_mullo_epi32(offsets_0, vksub);
+
+    // accumulators of partial sums
+    __m256 partialSum;
+
+    // load 8 uint8 values
+    const __m128i mm1 = _mm_loadu_si64((const __m128i_u*)code);
+    {
+        // convert uint8 values (low part of __m128i) to int32
+        // values
+        const __m256i idx1 = _mm256_cvtepu8_epi32(mm1);
 
-    // extract high and low __m128 regs from __m256
-    const __m128 h2 = _mm256_extractf128_ps(h1, 1);
-    const __m128 h3 = _mm256_castps256_ps128(h1);
+        // add offsets
+        const __m256i indices_to_read_from = _mm256_add_epi32(idx1, offsets_0);
 
-    // get a final hsum into all 4 regs
-    const __m128 h4 = _mm_add_ss(h2, h3);
+        // gather 8 values, similar to 8 operations of tab[idx]
+        __m256 collected =
+                _mm256_i32gather_ps(tab, indices_to_read_from, sizeof(float));
 
-    // extract f[0] from __m128
-    const float hsum = _mm_cvtss_f32(h4);
-    return hsum;
+        // collect partial sums
+        partialSum = collected;
+    }
+
+    // horizontal sum for partialSum
+    result = horizontal_sum(partialSum);
+    return result;
+}
+
+// processes four codes for M=4, ksub=256, nbits=8
+inline void distance_four_codes_avx2_pqdecoder8_m4(
+        // precomputed distances, layout (4, 256)
+        const float* sim_table,
+        // codes
+        const uint8_t* __restrict code0,
+        const uint8_t* __restrict code1,
+        const uint8_t* __restrict code2,
+        const uint8_t* __restrict code3,
+        // computed distances
+        float& result0,
+        float& result1,
+        float& result2,
+        float& result3) {
+    constexpr intptr_t N = 4;
+
+    const float* tab = sim_table;
+    constexpr size_t ksub = 1 << 8;
+
+    // process 8 values
+    const __m128i vksub = _mm_set1_epi32(ksub);
+    __m128i offsets_0 = _mm_setr_epi32(0, 1, 2, 3);
+    offsets_0 = _mm_mullo_epi32(offsets_0, vksub);
+
+    // accumulators of partial sums
+    __m128 partialSums[N];
+
+    // load 4 uint8 values
+    __m128i mm1[N];
+    mm1[0] = _mm_cvtsi32_si128(*((const int32_t*)code0));
+    mm1[1] = _mm_cvtsi32_si128(*((const int32_t*)code1));
+    mm1[2] = _mm_cvtsi32_si128(*((const int32_t*)code2));
+    mm1[3] = _mm_cvtsi32_si128(*((const int32_t*)code3));
+
+    for (intptr_t j = 0; j < N; j++) {
+        // convert uint8 values (low part of __m128i) to int32
+        // values
+        const __m128i idx1 = _mm_cvtepu8_epi32(mm1[j]);
+
+        // add offsets
+        const __m128i indices_to_read_from = _mm_add_epi32(idx1, offsets_0);
+
+        // gather 4 values, similar to 4 operations of tab[idx]
+        __m128 collected =
+                _mm_i32gather_ps(tab, indices_to_read_from, sizeof(float));
+
+        // collect partial sums
+        partialSums[j] = collected;
+    }
+
+    // horizontal sum for partialSum
+    result0 = horizontal_sum(partialSums[0]);
+    result1 = horizontal_sum(partialSums[1]);
+    result2 = horizontal_sum(partialSums[2]);
+    result3 = horizontal_sum(partialSums[3]);
+}
+
+// processes four codes for M=8, ksub=256, nbits=8
+inline void distance_four_codes_avx2_pqdecoder8_m8(
+        // precomputed distances, layout (8, 256)
+        const float* sim_table,
+        // codes
+        const uint8_t* __restrict code0,
+        const uint8_t* __restrict code1,
+        const uint8_t* __restrict code2,
+        const uint8_t* __restrict code3,
+        // computed distances
+        float& result0,
+        float& result1,
+        float& result2,
+        float& result3) {
+    constexpr intptr_t N = 4;
+
+    const float* tab = sim_table;
+    constexpr size_t ksub = 1 << 8;
+
+    // process 8 values
+    const __m256i vksub = _mm256_set1_epi32(ksub);
+    __m256i offsets_0 = _mm256_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7);
+    offsets_0 = _mm256_mullo_epi32(offsets_0, vksub);
+
+    // accumulators of partial sums
+    __m256 partialSums[N];
+
+    // load 8 uint8 values
+    __m128i mm1[N];
+    mm1[0] = _mm_loadu_si64((const __m128i_u*)code0);
+    mm1[1] = _mm_loadu_si64((const __m128i_u*)code1);
+    mm1[2] = _mm_loadu_si64((const __m128i_u*)code2);
+    mm1[3] = _mm_loadu_si64((const __m128i_u*)code3);
+
+    for (intptr_t j = 0; j < N; j++) {
+        // convert uint8 values (low part of __m128i) to int32
+        // values
+        const __m256i idx1 = _mm256_cvtepu8_epi32(mm1[j]);
+
+        // add offsets
+        const __m256i indices_to_read_from = _mm256_add_epi32(idx1, offsets_0);
+
+        // gather 8 values, similar to 8 operations of tab[idx]
+        __m256 collected =
+                _mm256_i32gather_ps(tab, indices_to_read_from, sizeof(float));
+
+        // collect partial sums
+        partialSums[j] = collected;
+    }
+
+    // horizontal sum for partialSum
+    result0 = horizontal_sum(partialSums[0]);
+    result1 = horizontal_sum(partialSums[1]);
+    result2 = horizontal_sum(partialSums[2]);
+    result3 = horizontal_sum(partialSums[3]);
 }
 
 } // namespace
@@ -41,36 +239,48 @@ namespace faiss {
 template <typename PQDecoderT>
 typename std::enable_if<!std::is_same<PQDecoderT, PQDecoder8>::value, float>::
         type inline distance_single_code_avx2(
-                // the product quantizer
-                const ProductQuantizer& pq,
+                // number of subquantizers
+                const size_t M,
+                // number of bits per quantization index
+                const size_t nbits,
                 // precomputed distances, layout (M, ksub)
                 const float* sim_table,
                 const uint8_t* code) {
     // default implementation
-    return distance_single_code_generic<PQDecoderT>(pq, sim_table, code);
+    return distance_single_code_generic<PQDecoderT>(M, nbits, sim_table, code);
 }
 
 template <typename PQDecoderT>
 typename std::enable_if<std::is_same<PQDecoderT, PQDecoder8>::value, float>::
         type inline distance_single_code_avx2(
-                // the product quantizer
-                const ProductQuantizer& pq,
+                // number of subquantizers
+                const size_t M,
+                // number of bits per quantization index
+                const size_t nbits,
                 // precomputed distances, layout (M, ksub)
                 const float* sim_table,
                 const uint8_t* code) {
+    if (M == 4) {
+        return distance_single_code_avx2_pqdecoder8_m4(sim_table, code);
+    }
+    if (M == 8) {
+        return distance_single_code_avx2_pqdecoder8_m8(sim_table, code);
+    }
+
     float result = 0;
+    constexpr size_t ksub = 1 << 8;
 
     size_t m = 0;
-    const size_t pqM16 = pq.M / 16;
+    const size_t pqM16 = M / 16;
 
     const float* tab = sim_table;
 
     if (pqM16 > 0) {
         // process 16 values per loop
 
-        const __m256i ksub = _mm256_set1_epi32(pq.ksub);
+        const __m256i vksub = _mm256_set1_epi32(ksub);
         __m256i offsets_0 = _mm256_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7);
-        offsets_0 = _mm256_mullo_epi32(offsets_0, ksub);
+        offsets_0 = _mm256_mullo_epi32(offsets_0, vksub);
 
         // accumulators of partial sums
         __m256 partialSum = _mm256_setzero_ps();
@@ -91,7 +301,7 @@ typename std::enable_if<std::is_same<PQDecoderT, PQDecoder8>::value, float>::
                 // gather 8 values, similar to 8 operations of tab[idx]
                 __m256 collected = _mm256_i32gather_ps(
                         tab, indices_to_read_from, sizeof(float));
-                tab += pq.ksub * 8;
+                tab += ksub * 8;
 
                 // collect partial sums
                 partialSum = _mm256_add_ps(partialSum, collected);
@@ -111,7 +321,7 @@ typename std::enable_if<std::is_same<PQDecoderT, PQDecoder8>::value, float>::
                 // gather 8 values, similar to 8 operations of tab[idx]
                 __m256 collected = _mm256_i32gather_ps(
                         tab, indices_to_read_from, sizeof(float));
-                tab += pq.ksub * 8;
+                tab += ksub * 8;
 
                 // collect partial sums
                 partialSum = _mm256_add_ps(partialSum, collected);
@@ -123,13 +333,13 @@ typename std::enable_if<std::is_same<PQDecoderT, PQDecoder8>::value, float>::
     }
 
     //
-    if (m < pq.M) {
+    if (m < M) {
         // process leftovers
-        PQDecoder8 decoder(code + m, pq.nbits);
+        PQDecoder8 decoder(code + m, nbits);
 
-        for (; m < pq.M; m++) {
+        for (; m < M; m++) {
             result += tab[decoder.decode()];
-            tab += pq.ksub;
+            tab += ksub;
         }
     }
 
@@ -140,8 +350,10 @@ template <typename PQDecoderT>
 typename std::enable_if<!std::is_same<PQDecoderT, PQDecoder8>::value, void>::
         type
         distance_four_codes_avx2(
-                // the product quantizer
-                const ProductQuantizer& pq,
+                // number of subquantizers
+                const size_t M,
+                // number of bits per quantization index
+                const size_t nbits,
                 // precomputed distances, layout (M, ksub)
                 const float* sim_table,
                 // codes
@@ -155,7 +367,8 @@ typename std::enable_if<!std::is_same<PQDecoderT, PQDecoder8>::value, void>::
                 float& result2,
                 float& result3) {
     distance_four_codes_generic<PQDecoderT>(
-            pq,
+            M,
+            nbits,
             sim_table,
             code0,
             code1,
@@ -171,8 +384,10 @@ typename std::enable_if<!std::is_same<PQDecoderT, PQDecoder8>::value, void>::
 template <typename PQDecoderT>
 typename std::enable_if<std::is_same<PQDecoderT, PQDecoder8>::value, void>::type
 distance_four_codes_avx2(
-        // the product quantizer
-        const ProductQuantizer& pq,
+        // number of subquantizers
+        const size_t M,
+        // number of bits per quantization index
+        const size_t nbits,
         // precomputed distances, layout (M, ksub)
         const float* sim_table,
         // codes
@@ -185,13 +400,41 @@ distance_four_codes_avx2(
         float& result1,
         float& result2,
         float& result3) {
+    if (M == 4) {
+        distance_four_codes_avx2_pqdecoder8_m4(
+                sim_table,
+                code0,
+                code1,
+                code2,
+                code3,
+                result0,
+                result1,
+                result2,
+                result3);
+        return;
+    }
+    if (M == 8) {
+        distance_four_codes_avx2_pqdecoder8_m8(
+                sim_table,
+                code0,
+                code1,
+                code2,
+                code3,
+                result0,
+                result1,
+                result2,
+                result3);
+        return;
+    }
+
     result0 = 0;
     result1 = 0;
     result2 = 0;
     result3 = 0;
+    constexpr size_t ksub = 1 << 8;
 
     size_t m = 0;
-    const size_t pqM16 = pq.M / 16;
+    const size_t pqM16 = M / 16;
 
     constexpr intptr_t N = 4;
 
@@ -199,9 +442,9 @@ distance_four_codes_avx2(
 
     if (pqM16 > 0) {
         // process 16 values per loop
-        const __m256i ksub = _mm256_set1_epi32(pq.ksub);
+        const __m256i vksub = _mm256_set1_epi32(ksub);
         __m256i offsets_0 = _mm256_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7);
-        offsets_0 = _mm256_mullo_epi32(offsets_0, ksub);
+        offsets_0 = _mm256_mullo_epi32(offsets_0, vksub);
 
         // accumulators of partial sums
         __m256 partialSums[N];
@@ -235,7 +478,7 @@ distance_four_codes_avx2(
                 // collect partial sums
                 partialSums[j] = _mm256_add_ps(partialSums[j], collected);
             }
-            tab += pq.ksub * 8;
+            tab += ksub * 8;
 
             // process next 8 codes
             for (intptr_t j = 0; j < N; j++) {
@@ -259,7 +502,7 @@ distance_four_codes_avx2(
                 partialSums[j] = _mm256_add_ps(partialSums[j], collected);
             }
 
-            tab += pq.ksub * 8;
+            tab += ksub * 8;
         }
 
         // horizontal sum for partialSum
@@ -270,18 +513,18 @@ distance_four_codes_avx2(
     }
 
     //
-    if (m < pq.M) {
+    if (m < M) {
         // process leftovers
-        PQDecoder8 decoder0(code0 + m, pq.nbits);
-        PQDecoder8 decoder1(code1 + m, pq.nbits);
-        PQDecoder8 decoder2(code2 + m, pq.nbits);
-        PQDecoder8 decoder3(code3 + m, pq.nbits);
-        for (; m < pq.M; m++) {
+        PQDecoder8 decoder0(code0 + m, nbits);
+        PQDecoder8 decoder1(code1 + m, nbits);
+        PQDecoder8 decoder2(code2 + m, nbits);
+        PQDecoder8 decoder3(code3 + m, nbits);
+        for (; m < M; m++) {
             result0 += tab[decoder0.decode()];
             result1 += tab[decoder1.decode()];
             result2 += tab[decoder2.decode()];
             result3 += tab[decoder3.decode()];
-            tab += pq.ksub;
+            tab += ksub;
         }
     }
 }