faiss 0.3.0 → 0.3.1

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

@@ -559,15 +559,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 +687,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 +703,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 +864,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 +955,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 +1049,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 +1067,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 +1164,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 +1296,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 +1314,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],

data/vendor/faiss/faiss/utils/sorting.cpp

@@ -123,7 +123,7 @@ void parallel_merge(
     }
 }
 
-}; // namespace
+} // namespace
 
 void fvec_argsort(size_t n, const float* vals, size_t* perm) {
     for (size_t i = 0; i < n; i++) {
@@ -544,7 +544,6 @@ void bucket_sort_inplace_parallel(
 
         // in this loop, we write elements collected in the previous round
         // and collect the elements that are overwritten for the next round
-        size_t tot_written = 0;
         int round = 0;
         for (;;) {
 #pragma omp barrier
@@ -554,9 +553,6 @@ void bucket_sort_inplace_parallel(
                 n_to_write += to_write_2.lims.back();
             }
 
-            tot_written += n_to_write;
-            // assert(tot_written <= nval);
-
 #pragma omp master
             {
                 if (verbose >= 1) {
@@ -689,4 +685,143 @@ void matrix_bucket_sort_inplace(
     }
 }
 
+/** Hashtable implementation for int64 -> int64 with external storage
+ * implemented for speed and parallel processing.
+ */
+
+namespace {
+
+int log2_capacity_to_log2_nbucket(int log2_capacity) {
+    return log2_capacity < 12    ? 0
+            : log2_capacity < 20 ? log2_capacity - 12
+                                 : 10;
+}
+
+// https://bigprimes.org/
+int64_t bigprime = 8955327411143;
+
+inline int64_t hash_function(int64_t x) {
+    return (x * 1000003) % bigprime;
+}
+
+} // anonymous namespace
+
+void hashtable_int64_to_int64_init(int log2_capacity, int64_t* tab) {
+    size_t capacity = (size_t)1 << log2_capacity;
+#pragma omp parallel for
+    for (int64_t i = 0; i < capacity; i++) {
+        tab[2 * i] = -1;
+        tab[2 * i + 1] = -1;
+    }
+}
+
+void hashtable_int64_to_int64_add(
+        int log2_capacity,
+        int64_t* tab,
+        size_t n,
+        const int64_t* keys,
+        const int64_t* vals) {
+    size_t capacity = (size_t)1 << log2_capacity;
+    std::vector<int64_t> hk(n);
+    std::vector<uint64_t> bucket_no(n);
+    int64_t mask = capacity - 1;
+    int log2_nbucket = log2_capacity_to_log2_nbucket(log2_capacity);
+    size_t nbucket = (size_t)1 << log2_nbucket;
+
+#pragma omp parallel for
+    for (int64_t i = 0; i < n; i++) {
+        hk[i] = hash_function(keys[i]) & mask;
+        bucket_no[i] = hk[i] >> (log2_capacity - log2_nbucket);
+    }
+
+    std::vector<int64_t> lims(nbucket + 1);
+    std::vector<int64_t> perm(n);
+    bucket_sort(
+            n,
+            bucket_no.data(),
+            nbucket,
+            lims.data(),
+            perm.data(),
+            omp_get_max_threads());
+
+    int num_errors = 0;
+#pragma omp parallel for reduction(+ : num_errors)
+    for (int64_t bucket = 0; bucket < nbucket; bucket++) {
+        size_t k0 = bucket << (log2_capacity - log2_nbucket);
+        size_t k1 = (bucket + 1) << (log2_capacity - log2_nbucket);
+
+        for (size_t i = lims[bucket]; i < lims[bucket + 1]; i++) {
+            int64_t j = perm[i];
+            assert(bucket_no[j] == bucket);
+            assert(hk[j] >= k0 && hk[j] < k1);
+            size_t slot = hk[j];
+            for (;;) {
+                if (tab[slot * 2] == -1) { // found!
+                    tab[slot * 2] = keys[j];
+                    tab[slot * 2 + 1] = vals[j];
+                    break;
+                } else if (tab[slot * 2] == keys[j]) { // overwrite!
+                    tab[slot * 2 + 1] = vals[j];
+                    break;
+                }
+                slot++;
+                if (slot == k1) {
+                    slot = k0;
+                }
+                if (slot == hk[j]) { // no free slot left in bucket
+                    num_errors++;
+                    break;
+                }
+            }
+            if (num_errors > 0) {
+                break;
+            }
+        }
+    }
+    FAISS_THROW_IF_NOT_MSG(num_errors == 0, "hashtable capacity exhausted");
+}
+
+void hashtable_int64_to_int64_lookup(
+        int log2_capacity,
+        const int64_t* tab,
+        size_t n,
+        const int64_t* keys,
+        int64_t* vals) {
+    size_t capacity = (size_t)1 << log2_capacity;
+    std::vector<int64_t> hk(n), bucket_no(n);
+    int64_t mask = capacity - 1;
+    int log2_nbucket = log2_capacity_to_log2_nbucket(log2_capacity);
+
+#pragma omp parallel for
+    for (int64_t i = 0; i < n; i++) {
+        int64_t k = keys[i];
+        int64_t hk = hash_function(k) & mask;
+        size_t slot = hk;
+
+        if (tab[2 * slot] == -1) { // not in table
+            vals[i] = -1;
+        } else if (tab[2 * slot] == k) { // found!
+            vals[i] = tab[2 * slot + 1];
+        } else { // need to search in [k0, k1)
+            size_t bucket = hk >> (log2_capacity - log2_nbucket);
+            size_t k0 = bucket << (log2_capacity - log2_nbucket);
+            size_t k1 = (bucket + 1) << (log2_capacity - log2_nbucket);
+            for (;;) {
+                if (tab[slot * 2] == k) { // found!
+                    vals[i] = tab[2 * slot + 1];
+                    break;
+                }
+                slot++;
+                if (slot == k1) {
+                    slot = k0;
+                }
+                if (slot == hk) { // bucket is full and not found
+                    vals[i] = -1;
+                    break;
+                }
+            }
+        }
+    }
+}
+
 } // namespace faiss

data/vendor/faiss/faiss/utils/sorting.h

@@ -68,4 +68,31 @@ void matrix_bucket_sort_inplace(
         int64_t* lims,
         int nt = 0);
 
+/** Hashtable implementation for int64 -> int64 with external storage
+ * implemented for fast batch add and lookup.
+ *
+ * tab is of size  2 * (1 << log2_capacity)
+ * n is the number of elements to add or search
+ *
+ * adding several values in a same batch: an arbitrary one gets added
+ * in different batches: the newer batch overwrites.
+ * raises an exception if capacity is exhausted.
+ */
+
+void hashtable_int64_to_int64_init(int log2_capacity, int64_t* tab);
+
+void hashtable_int64_to_int64_add(
+        int log2_capacity,
+        int64_t* tab,
+        size_t n,
+        const int64_t* keys,
+        const int64_t* vals);
+
+void hashtable_int64_to_int64_lookup(
+        int log2_capacity,
+        const int64_t* tab,
+        size_t n,
+        const int64_t* keys,
+        int64_t* vals);
+
 } // namespace faiss

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

@@ -28,6 +28,8 @@
 #include <omp.h>
 
 #include <algorithm>
+#include <set>
+#include <type_traits>
 #include <vector>
 
 #include <faiss/impl/AuxIndexStructures.h>
@@ -101,6 +103,9 @@ int sgemv_(
 
 namespace faiss {
 
+// this will be set at load time from GPU Faiss
+std::string gpu_compile_options;
+
 std::string get_compile_options() {
     std::string options;
 
@@ -110,13 +115,17 @@ std::string get_compile_options() {
 #endif
 
 #ifdef __AVX2__
-    options += "AVX2";
+    options += "AVX2 ";
+#elif __AVX512F__
+    options += "AVX512 ";
 #elif defined(__aarch64__)
-    options += "NEON";
+    options += "NEON ";
 #else
-    options += "GENERIC";
+    options += "GENERIC ";
 #endif
 
+    options += gpu_compile_options;
+
     return options;
 }
 
@@ -423,15 +432,35 @@ void bincode_hist(size_t n, size_t nbits, const uint8_t* codes, int* hist) {
     }
 }
 
-size_t ivec_checksum(size_t n, const int32_t* asigned) {
-    const uint32_t* a = reinterpret_cast<const uint32_t*>(asigned);
-    size_t cs = 112909;
+uint64_t ivec_checksum(size_t n, const int32_t* assigned) {
+    const uint32_t* a = reinterpret_cast<const uint32_t*>(assigned);
+    uint64_t cs = 112909;
     while (n--) {
         cs = cs * 65713 + a[n] * 1686049;
     }
     return cs;
 }
 
+uint64_t bvec_checksum(size_t n, const uint8_t* a) {
+    uint64_t cs = ivec_checksum(n / 4, (const int32_t*)a);
+    for (size_t i = n / 4 * 4; i < n; i++) {
+        cs = cs * 65713 + a[n] * 1686049;
+    }
+    return cs;
+}
+
+void bvecs_checksum(size_t n, size_t d, const uint8_t* a, uint64_t* cs) {
+    // MSVC can't accept unsigned index for #pragma omp parallel for
+    // so below codes only accept n <= std::numeric_limits<ssize_t>::max()
+    using ssize_t = std::make_signed<std::size_t>::type;
+    const ssize_t size = n;
+#pragma omp parallel for if (size > 1000)
+    for (ssize_t i_ = 0; i_ < size; i_++) {
+        const auto i = static_cast<std::size_t>(i_);
+        cs[i] = bvec_checksum(d, a + i * d);
+    }
+}
+
 const float* fvecs_maybe_subsample(
         size_t d,
         size_t* n,
@@ -528,4 +557,81 @@ bool check_openmp() {
     return true;
 }
 
+namespace {
+
+template <typename T>
+int64_t count_lt(int64_t n, const T* row, T threshold) {
+    for (int64_t i = 0; i < n; i++) {
+        if (!(row[i] < threshold)) {
+            return i;
+        }
+    }
+    return n;
+}
+
+template <typename T>
+int64_t count_gt(int64_t n, const T* row, T threshold) {
+    for (int64_t i = 0; i < n; i++) {
+        if (!(row[i] > threshold)) {
+            return i;
+        }
+    }
+    return n;
+}
+
+} // namespace
+
+template <typename T>
+void CombinerRangeKNN<T>::compute_sizes(int64_t* L_res_2) {
+    this->L_res = L_res_2;
+    L_res_2[0] = 0;
+    int64_t j = 0;
+    for (int64_t i = 0; i < nq; i++) {
+        int64_t n_in;
+        if (!mask || !mask[i]) {
+            const T* row = D + i * k;
+            n_in = keep_max ? count_gt(k, row, r2) : count_lt(k, row, r2);
+        } else {
+            n_in = lim_remain[j + 1] - lim_remain[j];
+            j++;
+        }
+        L_res_2[i + 1] = n_in; // L_res_2[i] + n_in;
+    }
+    // cumsum
+    for (int64_t i = 0; i < nq; i++) {
+        L_res_2[i + 1] += L_res_2[i];
+    }
+}
+
+template <typename T>
+void CombinerRangeKNN<T>::write_result(T* D_res, int64_t* I_res) {
+    FAISS_THROW_IF_NOT(L_res);
+    int64_t j = 0;
+    for (int64_t i = 0; i < nq; i++) {
+        int64_t n_in = L_res[i + 1] - L_res[i];
+        T* D_row = D_res + L_res[i];
+        int64_t* I_row = I_res + L_res[i];
+        if (!mask || !mask[i]) {
+            memcpy(D_row, D + i * k, n_in * sizeof(*D_row));
+            memcpy(I_row, I + i * k, n_in * sizeof(*I_row));
+        } else {
+            memcpy(D_row, D_remain + lim_remain[j], n_in * sizeof(*D_row));
+            memcpy(I_row, I_remain + lim_remain[j], n_in * sizeof(*I_row));
+            j++;
+        }
+    }
+}
+
+// explicit template instantiations
+template struct CombinerRangeKNN<float>;
+template struct CombinerRangeKNN<int16_t>;
+
+void CodeSet::insert(size_t n, const uint8_t* codes, bool* inserted) {
+    for (size_t i = 0; i < n; i++) {
+        auto res = s.insert(
+                std::vector<uint8_t>(codes + i * d, codes + i * d + d));
+        inserted[i] = res.second;
+    }
+}
+
 } // namespace faiss