faiss 0.2.7 → 0.3.1

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

@@ -0,0 +1,29 @@
+/**
+ * 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 <arm_neon.h>
+#include <cstdint>
+
+namespace faiss {
+
+inline uint16_t encode_fp16(float x) {
+    float32x4_t fx4 = vdupq_n_f32(x);
+    float16x4_t f16x4 = vcvt_f16_f32(fx4);
+    uint16x4_t ui16x4 = vreinterpret_u16_f16(f16x4);
+    return vduph_lane_u16(ui16x4, 3);
+}
+
+inline float decode_fp16(uint16_t x) {
+    uint16x4_t ui16x4 = vdup_n_u16(x);
+    float16x4_t f16x4 = vreinterpret_f16_u16(ui16x4);
+    float32x4_t fx4 = vcvt_f32_f16(f16x4);
+    return vdups_laneq_f32(fx4, 3);
+}
+
+} // namespace faiss

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

@@ -13,6 +13,8 @@
 
 #if defined(__F16C__)
 #include <faiss/utils/fp16-fp16c.h>
+#elif defined(__aarch64__)
+#include <faiss/utils/fp16-arm.h>
 #else
 #include <faiss/utils/fp16-inl.h>
 #endif

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

@@ -5,14 +5,13 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-// -*- c++ -*-
-
 /*
  * Implementation of Hamming related functions (distances, smallest distance
  * selection with regular heap|radix and probabilistic heap|radix.
  *
  * IMPLEMENTATION NOTES
- * Bitvectors are generally assumed to be multiples of 64 bits.
+ * Optimal speed is typically obtained for vector sizes of multiples of 64
+ * bits.
  *
  * hamdis_t is used for distances because at this time
  * it is not clear how we will need to balance
@@ -20,15 +19,13 @@
  * - memory usage
  * - cache-misses when dealing with large volumes of data (lower bits is better)
  *
- * The hamdis_t should optimally be compatibe with one of the Torch Storage
- * (Byte,Short,Long) and therefore should be signed for 2-bytes and 4-bytes
  */
 
 #include <faiss/utils/hamming.h>
 
-#include <math.h>
-#include <stdio.h>
 #include <algorithm>
+#include <cmath>
+#include <cstdio>
 #include <memory>
 #include <vector>
 
@@ -38,8 +35,6 @@
 #include <faiss/utils/approx_topk_hamming/approx_topk_hamming.h>
 #include <faiss/utils/utils.h>
 
-static const size_t BLOCKSIZE_QUERY = 8192;
-
 namespace faiss {
 
 size_t hamming_batch_size = 65536;
@@ -165,9 +160,11 @@ size_t match_hamming_thres(
     return posm;
 }
 
+namespace {
+
 /* Return closest neighbors w.r.t Hamming distance, using a heap. */
 template <class HammingComputer>
-static void hammings_knn_hc(
+void hammings_knn_hc(
         int bytes_per_code,
         int_maxheap_array_t* __restrict ha,
         const uint8_t* __restrict bs1,
@@ -234,7 +231,7 @@ static void hammings_knn_hc(
 
 /* Return closest neighbors w.r.t Hamming distance, using max count. */
 template <class HammingComputer>
-static void hammings_knn_mc(
+void hammings_knn_mc(
         int bytes_per_code,
         const uint8_t* __restrict a,
         const uint8_t* __restrict b,
@@ -272,10 +269,10 @@ static void hammings_knn_mc(
         HCounterState<HammingComputer>& csi = cs[i];
 
         int nres = 0;
-        for (int b = 0; b < nBuckets && nres < k; b++) {
-            for (int l = 0; l < csi.counters[b] && nres < k; l++) {
-                labels[i * k + nres] = csi.ids_per_dis[b * k + l];
-                distances[i * k + nres] = b;
+        for (int b_2 = 0; b_2 < nBuckets && nres < k; b_2++) {
+            for (int l = 0; l < csi.counters[b_2] && nres < k; l++) {
+                labels[i * k + nres] = csi.ids_per_dis[b_2 * k + l];
+                distances[i * k + nres] = b_2;
                 nres++;
             }
         }
@@ -287,6 +284,63 @@ static void hammings_knn_mc(
     }
 }
 
+template <class HammingComputer>
+void hamming_range_search(
+        const uint8_t* a,
+        const uint8_t* b,
+        size_t na,
+        size_t nb,
+        int radius,
+        size_t code_size,
+        RangeSearchResult* res) {
+#pragma omp parallel
+    {
+        RangeSearchPartialResult pres(res);
+
+#pragma omp for
+        for (int64_t i = 0; i < na; i++) {
+            HammingComputer hc(a + i * code_size, code_size);
+            const uint8_t* yi = b;
+            RangeQueryResult& qres = pres.new_result(i);
+
+            for (size_t j = 0; j < nb; j++) {
+                int dis = hc.hamming(yi);
+                if (dis < radius) {
+                    qres.add(dis, j);
+                }
+                yi += code_size;
+            }
+        }
+        pres.finalize();
+    }
+}
+
+struct Run_hammings_knn_hc {
+    using T = void;
+    template <class HammingComputer, class... Types>
+    void f(Types... args) {
+        hammings_knn_hc<HammingComputer>(args...);
+    }
+};
+
+struct Run_hammings_knn_mc {
+    using T = void;
+    template <class HammingComputer, class... Types>
+    void f(Types... args) {
+        hammings_knn_mc<HammingComputer>(args...);
+    }
+};
+
+struct Run_hamming_range_search {
+    using T = void;
+    template <class HammingComputer, class... Types>
+    void f(Types... args) {
+        hamming_range_search<HammingComputer>(args...);
+    }
+};
+
+} // namespace
+
 /* Functions to maps vectors to bits. Assume proper allocation done beforehand,
    meaning that b should be be able to receive as many bits as x may produce. */
 
@@ -437,28 +491,9 @@ void hammings_knn_hc(
         size_t ncodes,
         int order,
         ApproxTopK_mode_t approx_topk_mode) {
-    switch (ncodes) {
-        case 4:
-            hammings_knn_hc<faiss::HammingComputer4>(
-                    4, ha, a, b, nb, order, true, approx_topk_mode);
-            break;
-        case 8:
-            hammings_knn_hc<faiss::HammingComputer8>(
-                    8, ha, a, b, nb, order, true, approx_topk_mode);
-            break;
-        case 16:
-            hammings_knn_hc<faiss::HammingComputer16>(
-                    16, ha, a, b, nb, order, true, approx_topk_mode);
-            break;
-        case 32:
-            hammings_knn_hc<faiss::HammingComputer32>(
-                    32, ha, a, b, nb, order, true, approx_topk_mode);
-            break;
-        default:
-            hammings_knn_hc<faiss::HammingComputerDefault>(
-                    ncodes, ha, a, b, nb, order, true, approx_topk_mode);
-            break;
-    }
+    Run_hammings_knn_hc r;
+    dispatch_HammingComputer(
+            ncodes, r, ncodes, ha, a, b, nb, order, true, approx_topk_mode);
 }
 
 void hammings_knn_mc(
@@ -470,58 +505,9 @@ void hammings_knn_mc(
         size_t ncodes,
         int32_t* __restrict distances,
         int64_t* __restrict labels) {
-    switch (ncodes) {
-        case 4:
-            hammings_knn_mc<faiss::HammingComputer4>(
-                    4, a, b, na, nb, k, distances, labels);
-            break;
-        case 8:
-            hammings_knn_mc<faiss::HammingComputer8>(
-                    8, a, b, na, nb, k, distances, labels);
-            break;
-        case 16:
-            hammings_knn_mc<faiss::HammingComputer16>(
-                    16, a, b, na, nb, k, distances, labels);
-            break;
-        case 32:
-            hammings_knn_mc<faiss::HammingComputer32>(
-                    32, a, b, na, nb, k, distances, labels);
-            break;
-        default:
-            hammings_knn_mc<faiss::HammingComputerDefault>(
-                    ncodes, a, b, na, nb, k, distances, labels);
-            break;
-    }
-}
-template <class HammingComputer>
-static void hamming_range_search_template(
-        const uint8_t* a,
-        const uint8_t* b,
-        size_t na,
-        size_t nb,
-        int radius,
-        size_t code_size,
-        RangeSearchResult* res) {
-#pragma omp parallel
-    {
-        RangeSearchPartialResult pres(res);
-
-#pragma omp for
-        for (int64_t i = 0; i < na; i++) {
-            HammingComputer hc(a + i * code_size, code_size);
-            const uint8_t* yi = b;
-            RangeQueryResult& qres = pres.new_result(i);
-
-            for (size_t j = 0; j < nb; j++) {
-                int dis = hc.hamming(yi);
-                if (dis < radius) {
-                    qres.add(dis, j);
-                }
-                yi += code_size;
-            }
-        }
-        pres.finalize();
-    }
+    Run_hammings_knn_mc r;
+    dispatch_HammingComputer(
+            ncodes, r, ncodes, a, b, na, nb, k, distances, labels);
 }
 
 void hamming_range_search(
@@ -532,27 +518,9 @@ void hamming_range_search(
         int radius,
         size_t code_size,
         RangeSearchResult* result) {
-#define HC(name) \
-    hamming_range_search_template<name>(a, b, na, nb, radius, code_size, result)
-
-    switch (code_size) {
-        case 4:
-            HC(HammingComputer4);
-            break;
-        case 8:
-            HC(HammingComputer8);
-            break;
-        case 16:
-            HC(HammingComputer16);
-            break;
-        case 32:
-            HC(HammingComputer32);
-            break;
-        default:
-            HC(HammingComputerDefault);
-            break;
-    }
-#undef HC
+    Run_hamming_range_search r;
+    dispatch_HammingComputer(
+            code_size, r, a, b, na, nb, radius, code_size, result);
 }
 
 /* Count number of matches given a max threshold            */
@@ -711,4 +679,88 @@ void generalized_hammings_knn_hc(
         ha->reorder();
 }
 
+void pack_bitstrings(
+        size_t n,
+        size_t M,
+        int nbit,
+        const int32_t* unpacked,
+        uint8_t* packed,
+        size_t code_size) {
+    FAISS_THROW_IF_NOT(code_size >= (M * nbit + 7) / 8);
+#pragma omp parallel for if (n > 1000)
+    for (int64_t i = 0; i < n; i++) {
+        const int32_t* in = unpacked + i * M;
+        uint8_t* out = packed + i * code_size;
+        BitstringWriter wr(out, code_size);
+        for (int j = 0; j < M; j++) {
+            wr.write(in[j], nbit);
+        }
+    }
+}
+
+void pack_bitstrings(
+        size_t n,
+        size_t M,
+        const int32_t* nbit,
+        const int32_t* unpacked,
+        uint8_t* packed,
+        size_t code_size) {
+    int totbit = 0;
+    for (int j = 0; j < M; j++) {
+        totbit += nbit[j];
+    }
+    FAISS_THROW_IF_NOT(code_size >= (totbit + 7) / 8);
+#pragma omp parallel for if (n > 1000)
+    for (int64_t i = 0; i < n; i++) {
+        const int32_t* in = unpacked + i * M;
+        uint8_t* out = packed + i * code_size;
+        BitstringWriter wr(out, code_size);
+        for (int j = 0; j < M; j++) {
+            wr.write(in[j], nbit[j]);
+        }
+    }
+}
+
+void unpack_bitstrings(
+        size_t n,
+        size_t M,
+        int nbit,
+        const uint8_t* packed,
+        size_t code_size,
+        int32_t* unpacked) {
+    FAISS_THROW_IF_NOT(code_size >= (M * nbit + 7) / 8);
+#pragma omp parallel for if (n > 1000)
+    for (int64_t i = 0; i < n; i++) {
+        const uint8_t* in = packed + i * code_size;
+        int32_t* out = unpacked + i * M;
+        BitstringReader rd(in, code_size);
+        for (int j = 0; j < M; j++) {
+            out[j] = rd.read(nbit);
+        }
+    }
+}
+
+void unpack_bitstrings(
+        size_t n,
+        size_t M,
+        const int32_t* nbit,
+        const uint8_t* packed,
+        size_t code_size,
+        int32_t* unpacked) {
+    int totbit = 0;
+    for (int j = 0; j < M; j++) {
+        totbit += nbit[j];
+    }
+    FAISS_THROW_IF_NOT(code_size >= (totbit + 7) / 8);
+#pragma omp parallel for if (n > 1000)
+    for (int64_t i = 0; i < n; i++) {
+        const uint8_t* in = packed + i * code_size;
+        int32_t* out = unpacked + i * M;
+        BitstringReader rd(in, code_size);
+        for (int j = 0; j < M; j++) {
+            out[j] = rd.read(nbit[j]);
+        }
+    }
+}
+
 } // namespace faiss

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

@@ -222,6 +222,64 @@ void generalized_hammings_knn_hc(
         size_t code_size,
         int ordered = true);
 
+/** Pack a set of n codes of size M * nbit
+ *
+ * @param n           number of codes to pack
+ * @param M           number of elementary codes per code
+ * @param nbit        number of bits per elementary code
+ * @param unpacked    input unpacked codes, size (n, M)
+ * @param packed      output packed codes, size (n, code_size)
+ * @param code_size   should be >= ceil(M * nbit / 8)
+ */
+void pack_bitstrings(
+        size_t n,
+        size_t M,
+        int nbit,
+        const int32_t* unpacked,
+        uint8_t* packed,
+        size_t code_size);
+
+/** Pack a set of n codes of variable sizes
+ *
+ * @param nbit       number of bits per entry (size M)
+ */
+void pack_bitstrings(
+        size_t n,
+        size_t M,
+        const int32_t* nbits,
+        const int32_t* unpacked,
+        uint8_t* packed,
+        size_t code_size);
+
+/** Unpack a set of n codes of size M * nbit
+ *
+ * @param n           number of codes to pack
+ * @param M           number of elementary codes per code
+ * @param nbit        number of bits per elementary code
+ * @param unpacked    input unpacked codes, size (n, M)
+ * @param packed      output packed codes, size (n, code_size)
+ * @param code_size   should be >= ceil(M * nbit / 8)
+ */
+void unpack_bitstrings(
+        size_t n,
+        size_t M,
+        int nbit,
+        const uint8_t* packed,
+        size_t code_size,
+        int32_t* unpacked);
+
+/** Unpack a set of n codes of variable sizes
+ *
+ * @param nbit       number of bits per entry (size M)
+ */
+void unpack_bitstrings(
+        size_t n,
+        size_t M,
+        const int32_t* nbits,
+        const uint8_t* packed,
+        size_t code_size,
+        int32_t* unpacked);
+
 } // namespace faiss
 
 #include <faiss/utils/hamming-inl.h>

data/vendor/faiss/faiss/utils/hamming_distance/avx2-inl.h

@@ -259,8 +259,8 @@ struct HammingComputerDefault {
         set(a8, code_size);
     }
 
-    void set(const uint8_t* a8, int code_size) {
-        this->a8 = a8;
+    void set(const uint8_t* a8_2, int code_size) {
+        this->a8 = a8_2;
         quotient8 = code_size / 8;
         remainder8 = code_size % 8;
     }
@@ -277,24 +277,31 @@ struct HammingComputerDefault {
                     len -= 8;
                     accu += popcount64(a64[i] ^ b64[i]);
                     i++;
+                    [[fallthrough]];
                     case 7:
                         accu += popcount64(a64[i] ^ b64[i]);
                         i++;
+                        [[fallthrough]];
                     case 6:
                         accu += popcount64(a64[i] ^ b64[i]);
                         i++;
+                        [[fallthrough]];
                     case 5:
                         accu += popcount64(a64[i] ^ b64[i]);
                         i++;
+                        [[fallthrough]];
                     case 4:
                         accu += popcount64(a64[i] ^ b64[i]);
                         i++;
+                        [[fallthrough]];
                     case 3:
                         accu += popcount64(a64[i] ^ b64[i]);
                         i++;
+                        [[fallthrough]];
                     case 2:
                         accu += popcount64(a64[i] ^ b64[i]);
                         i++;
+                        [[fallthrough]];
                     case 1:
                         accu += popcount64(a64[i] ^ b64[i]);
                         i++;
@@ -306,18 +313,25 @@ struct HammingComputerDefault {
             switch (remainder8) {
                 case 7:
                     accu += hamdis_tab_ham_bytes[a[6] ^ b[6]];
+                    [[fallthrough]];
                 case 6:
                     accu += hamdis_tab_ham_bytes[a[5] ^ b[5]];
+                    [[fallthrough]];
                 case 5:
                     accu += hamdis_tab_ham_bytes[a[4] ^ b[4]];
+                    [[fallthrough]];
                 case 4:
                     accu += hamdis_tab_ham_bytes[a[3] ^ b[3]];
+                    [[fallthrough]];
                 case 3:
                     accu += hamdis_tab_ham_bytes[a[2] ^ b[2]];
+                    [[fallthrough]];
                 case 2:
                     accu += hamdis_tab_ham_bytes[a[1] ^ b[1]];
+                    [[fallthrough]];
                 case 1:
                     accu += hamdis_tab_ham_bytes[a[0] ^ b[0]];
+                    [[fallthrough]];
                 default:
                     break;
             }
@@ -331,93 +345,6 @@ struct HammingComputerDefault {
     }
 };
 
-// more inefficient than HammingComputerDefault (obsolete)
-struct HammingComputerM8 {
-    const uint64_t* a;
-    int n;
-
-    HammingComputerM8() {}
-
-    HammingComputerM8(const uint8_t* a8, int code_size) {
-        set(a8, code_size);
-    }
-
-    void set(const uint8_t* a8, int code_size) {
-        assert(code_size % 8 == 0);
-        a = (uint64_t*)a8;
-        n = code_size / 8;
-    }
-
-    int hamming(const uint8_t* b8) const {
-        const uint64_t* b = (uint64_t*)b8;
-        int accu = 0;
-        for (int i = 0; i < n; i++)
-            accu += popcount64(a[i] ^ b[i]);
-        return accu;
-    }
-
-    inline int get_code_size() const {
-        return n * 8;
-    }
-};
-
-// more inefficient than HammingComputerDefault (obsolete)
-struct HammingComputerM4 {
-    const uint32_t* a;
-    int n;
-
-    HammingComputerM4() {}
-
-    HammingComputerM4(const uint8_t* a4, int code_size) {
-        set(a4, code_size);
-    }
-
-    void set(const uint8_t* a4, int code_size) {
-        assert(code_size % 4 == 0);
-        a = (uint32_t*)a4;
-        n = code_size / 4;
-    }
-
-    int hamming(const uint8_t* b8) const {
-        const uint32_t* b = (uint32_t*)b8;
-        int accu = 0;
-        for (int i = 0; i < n; i++)
-            accu += popcount64(a[i] ^ b[i]);
-        return accu;
-    }
-
-    inline int get_code_size() const {
-        return n * 4;
-    }
-};
-
-/***************************************************************************
- * Equivalence with a template class when code size is known at compile time
- **************************************************************************/
-
-// default template
-template <int CODE_SIZE>
-struct HammingComputer : HammingComputerDefault {
-    HammingComputer(const uint8_t* a, int code_size)
-            : HammingComputerDefault(a, code_size) {}
-};
-
-#define SPECIALIZED_HC(CODE_SIZE)                                    \
-    template <>                                                      \
-    struct HammingComputer<CODE_SIZE> : HammingComputer##CODE_SIZE { \
-        HammingComputer(const uint8_t* a)                            \
-                : HammingComputer##CODE_SIZE(a, CODE_SIZE) {}        \
-    }
-
-SPECIALIZED_HC(4);
-SPECIALIZED_HC(8);
-SPECIALIZED_HC(16);
-SPECIALIZED_HC(20);
-SPECIALIZED_HC(32);
-SPECIALIZED_HC(64);
-
-#undef SPECIALIZED_HC
-
 /***************************************************************************
  * generalized Hamming = number of bytes that are different between
  * two codes.

data/vendor/faiss/faiss/utils/hamming_distance/common.h

@@ -17,6 +17,7 @@ using hamdis_t = int32_t;
 
 namespace faiss {
 
+// trust the compiler to provide efficient popcount implementations
 inline int popcount32(uint32_t x) {
     return __builtin_popcount(x);
 }