faiss 0.3.1 → 0.3.2

data/vendor/faiss/faiss/utils/extra_distances-inl.h

@@ -9,7 +9,9 @@
  *  and inner product */
 
 #include <faiss/MetricType.h>
+#include <faiss/impl/FaissAssert.h>
 #include <faiss/utils/distances.h>
+#include <cmath>
 #include <type_traits>
 
 namespace faiss {
@@ -130,4 +132,72 @@ inline float VectorDistance<METRIC_Jaccard>::operator()(
     return accu_num / accu_den;
 }
 
+template <>
+inline float VectorDistance<METRIC_NaNEuclidean>::operator()(
+        const float* x,
+        const float* y) const {
+    // https://scikit-learn.org/stable/modules/generated/sklearn.metrics.pairwise.nan_euclidean_distances.html
+    float accu = 0;
+    size_t present = 0;
+    for (size_t i = 0; i < d; i++) {
+        if (!std::isnan(x[i]) && !std::isnan(y[i])) {
+            float diff = x[i] - y[i];
+            accu += diff * diff;
+            present++;
+        }
+    }
+    if (present == 0) {
+        return NAN;
+    }
+    return float(d) / float(present) * accu;
+}
+
+template <>
+inline float VectorDistance<METRIC_ABS_INNER_PRODUCT>::operator()(
+        const float* x,
+        const float* y) const {
+    float accu = 0;
+    for (size_t i = 0; i < d; i++) {
+        accu += fabs(x[i] * y[i]);
+    }
+    return accu;
+}
+
+/***************************************************************************
+ * Dispatching function that takes a metric type and a consumer object
+ * the consumer object should contain a retun type T and a operation template
+ * function f() that is called to perform the operation. The first argument
+ * of the function is the VectorDistance object. The rest are passed in as is.
+ **************************************************************************/
+
+template <class Consumer, class... Types>
+typename Consumer::T dispatch_VectorDistance(
+        size_t d,
+        MetricType metric,
+        float metric_arg,
+        Consumer& consumer,
+        Types... args) {
+    switch (metric) {
+#define DISPATCH_VD(mt)                                              \
+    case mt: {                                                       \
+        VectorDistance<mt> vd = {d, metric_arg};                     \
+        return consumer.template f<VectorDistance<mt>>(vd, args...); \
+    }
+        DISPATCH_VD(METRIC_INNER_PRODUCT);
+        DISPATCH_VD(METRIC_L2);
+        DISPATCH_VD(METRIC_L1);
+        DISPATCH_VD(METRIC_Linf);
+        DISPATCH_VD(METRIC_Lp);
+        DISPATCH_VD(METRIC_Canberra);
+        DISPATCH_VD(METRIC_BrayCurtis);
+        DISPATCH_VD(METRIC_JensenShannon);
+        DISPATCH_VD(METRIC_Jaccard);
+        DISPATCH_VD(METRIC_NaNEuclidean);
+        DISPATCH_VD(METRIC_ABS_INNER_PRODUCT);
+        default:
+            FAISS_THROW_FMT("Invalid metric %d", metric);
+    }
+#undef DISPATCH_VD
+}
+
 } // namespace faiss

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

@@ -26,73 +26,77 @@ namespace faiss {
 
 namespace {
 
-template <class VD>
-void pairwise_extra_distances_template(
-        VD vd,
-        int64_t nq,
-        const float* xq,
-        int64_t nb,
-        const float* xb,
-        float* dis,
-        int64_t ldq,
-        int64_t ldb,
-        int64_t ldd) {
+struct Run_pairwise_extra_distances {
+    using T = void;
+
+    template <class VD>
+    void f(VD vd,
+           int64_t nq,
+           const float* xq,
+           int64_t nb,
+           const float* xb,
+           float* dis,
+           int64_t ldq,
+           int64_t ldb,
+           int64_t ldd) {
 #pragma omp parallel for if (nq > 10)
-    for (int64_t i = 0; i < nq; i++) {
-        const float* xqi = xq + i * ldq;
-        const float* xbj = xb;
-        float* disi = dis + ldd * i;
-
-        for (int64_t j = 0; j < nb; j++) {
-            disi[j] = vd(xqi, xbj);
-            xbj += ldb;
+        for (int64_t i = 0; i < nq; i++) {
+            const float* xqi = xq + i * ldq;
+            const float* xbj = xb;
+            float* disi = dis + ldd * i;
+
+            for (int64_t j = 0; j < nb; j++) {
+                disi[j] = vd(xqi, xbj);
+                xbj += ldb;
+            }
         }
     }
-}
-
-template <class VD, class C>
-void knn_extra_metrics_template(
-        VD vd,
-        const float* x,
-        const float* y,
-        size_t nx,
-        size_t ny,
-        HeapArray<C>* res) {
-    size_t k = res->k;
-    size_t d = vd.d;
-    size_t check_period = InterruptCallback::get_period_hint(ny * d);
-    check_period *= omp_get_max_threads();
+};
 
-    for (size_t i0 = 0; i0 < nx; i0 += check_period) {
-        size_t i1 = std::min(i0 + check_period, nx);
+struct Run_knn_extra_metrics {
+    using T = void;
+    template <class VD>
+    void f(VD vd,
+           const float* x,
+           const float* y,
+           size_t nx,
+           size_t ny,
+           size_t k,
+           float* distances,
+           int64_t* labels) {
+        size_t d = vd.d;
+        using C = typename VD::C;
+        size_t check_period = InterruptCallback::get_period_hint(ny * d);
+        check_period *= omp_get_max_threads();
+
+        for (size_t i0 = 0; i0 < nx; i0 += check_period) {
+            size_t i1 = std::min(i0 + check_period, nx);
 
 #pragma omp parallel for
-        for (int64_t i = i0; i < i1; i++) {
-            const float* x_i = x + i * d;
-            const float* y_j = y;
-            size_t j;
-            float* simi = res->get_val(i);
-            int64_t* idxi = res->get_ids(i);
-
-            // maxheap_heapify(k, simi, idxi);
-            heap_heapify<C>(k, simi, idxi);
-            for (j = 0; j < ny; j++) {
-                float disij = vd(x_i, y_j);
-
-                // if (disij < simi[0]) {
-                if ((!vd.is_similarity && (disij < simi[0])) ||
-                    (vd.is_similarity && (disij > simi[0]))) {
-                    // maxheap_replace_top(k, simi, idxi, disij, j);
-                    heap_replace_top<C>(k, simi, idxi, disij, j);
+            for (int64_t i = i0; i < i1; i++) {
+                const float* x_i = x + i * d;
+                const float* y_j = y;
+                size_t j;
+                float* simi = distances + k * i;
+                int64_t* idxi = labels + k * i;
+
+                // maxheap_heapify(k, simi, idxi);
+                heap_heapify<C>(k, simi, idxi);
+                for (j = 0; j < ny; j++) {
+                    float disij = vd(x_i, y_j);
+
+                    if (C::cmp(simi[0], disij)) {
+                        heap_replace_top<C>(k, simi, idxi, disij, j);
+                    }
+                    y_j += d;
                 }
-                y_j += d;
+                // maxheap_reorder(k, simi, idxi);
+                heap_reorder<C>(k, simi, idxi);
             }
-            // maxheap_reorder(k, simi, idxi);
-            heap_reorder<C>(k, simi, idxi);
+            InterruptCallback::check();
         }
-        InterruptCallback::check();
     }
-}
+};
 
 template <class VD>
 struct ExtraDistanceComputer : FlatCodesDistanceComputer {
@@ -125,6 +129,19 @@ struct ExtraDistanceComputer : FlatCodesDistanceComputer {
     }
 };
 
+struct Run_get_distance_computer {
+    using T = FlatCodesDistanceComputer*;
+
+    template <class VD>
+    FlatCodesDistanceComputer* f(
+            VD vd,
+            const float* xb,
+            size_t nb,
+            const float* q = nullptr) {
+        return new ExtraDistanceComputer<VD>(vd, xb, nb, q);
+    }
+};
+
 } // anonymous namespace
 
 void pairwise_extra_distances(
@@ -148,29 +165,11 @@ void pairwise_extra_distances(
     if (ldd == -1)
         ldd = nb;
 
-    switch (mt) {
-#define HANDLE_VAR(kw)                                            \
-    case METRIC_##kw: {                                           \
-        VectorDistance<METRIC_##kw> vd = {(size_t)d, metric_arg}; \
-        pairwise_extra_distances_template(                        \
-                vd, nq, xq, nb, xb, dis, ldq, ldb, ldd);          \
-        break;                                                    \
-    }
-        HANDLE_VAR(L2);
-        HANDLE_VAR(L1);
-        HANDLE_VAR(Linf);
-        HANDLE_VAR(Canberra);
-        HANDLE_VAR(BrayCurtis);
-        HANDLE_VAR(JensenShannon);
-        HANDLE_VAR(Lp);
-        HANDLE_VAR(Jaccard);
-#undef HANDLE_VAR
-        default:
-            FAISS_THROW_MSG("metric type not implemented");
-    }
+    Run_pairwise_extra_distances run;
+    dispatch_VectorDistance(
+            d, mt, metric_arg, run, nq, xq, nb, xb, dis, ldq, ldb, ldd);
 }
 
-template <class C>
 void knn_extra_metrics(
         const float* x,
         const float* y,
@@ -179,73 +178,22 @@ void knn_extra_metrics(
         size_t ny,
         MetricType mt,
         float metric_arg,
-        HeapArray<C>* res) {
-    switch (mt) {
-#define HANDLE_VAR(kw)                                            \
-    case METRIC_##kw: {                                           \
-        VectorDistance<METRIC_##kw> vd = {(size_t)d, metric_arg}; \
-        knn_extra_metrics_template(vd, x, y, nx, ny, res);        \
-        break;                                                    \
-    }
-        HANDLE_VAR(L2);
-        HANDLE_VAR(L1);
-        HANDLE_VAR(Linf);
-        HANDLE_VAR(Canberra);
-        HANDLE_VAR(BrayCurtis);
-        HANDLE_VAR(JensenShannon);
-        HANDLE_VAR(Lp);
-        HANDLE_VAR(Jaccard);
-#undef HANDLE_VAR
-        default:
-            FAISS_THROW_MSG("metric type not implemented");
-    }
+        size_t k,
+        float* distances,
+        int64_t* indexes) {
+    Run_knn_extra_metrics run;
+    dispatch_VectorDistance(
+            d, mt, metric_arg, run, x, y, nx, ny, k, distances, indexes);
 }
 
-template void knn_extra_metrics<CMax<float, int64_t>>(
-        const float* x,
-        const float* y,
-        size_t d,
-        size_t nx,
-        size_t ny,
-        MetricType mt,
-        float metric_arg,
-        HeapArray<CMax<float, int64_t>>* res);
-
-template void knn_extra_metrics<CMin<float, int64_t>>(
-        const float* x,
-        const float* y,
-        size_t d,
-        size_t nx,
-        size_t ny,
-        MetricType mt,
-        float metric_arg,
-        HeapArray<CMin<float, int64_t>>* res);
-
 FlatCodesDistanceComputer* get_extra_distance_computer(
         size_t d,
         MetricType mt,
         float metric_arg,
         size_t nb,
         const float* xb) {
-    switch (mt) {
-#define HANDLE_VAR(kw)                                                 \
-    case METRIC_##kw: {                                                \
-        VectorDistance<METRIC_##kw> vd = {(size_t)d, metric_arg};      \
-        return new ExtraDistanceComputer<VectorDistance<METRIC_##kw>>( \
-                vd, xb, nb);                                           \
-    }
-        HANDLE_VAR(L2);
-        HANDLE_VAR(L1);
-        HANDLE_VAR(Linf);
-        HANDLE_VAR(Canberra);
-        HANDLE_VAR(BrayCurtis);
-        HANDLE_VAR(JensenShannon);
-        HANDLE_VAR(Lp);
-        HANDLE_VAR(Jaccard);
-#undef HANDLE_VAR
-        default:
-            FAISS_THROW_MSG("metric type not implemented");
-    }
+    Run_get_distance_computer run;
+    return dispatch_VectorDistance(d, mt, metric_arg, run, xb, nb);
 }
 
 } // namespace faiss

data/vendor/faiss/faiss/utils/extra_distances.h

@@ -33,7 +33,6 @@ void pairwise_extra_distances(
         int64_t ldb = -1,
         int64_t ldd = -1);
 
-template <class C>
 void knn_extra_metrics(
         const float* x,
         const float* y,
@@ -42,7 +41,9 @@ void knn_extra_metrics(
         size_t ny,
         MetricType mt,
         float metric_arg,
-        HeapArray<C>* res);
+        size_t k,
+        float* distances,
+        int64_t* indexes);
 
 /** get a DistanceComputer that refers to this type of distance and
  *  indexes a flat array of size nb */

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

@@ -364,7 +364,7 @@ void fvec2bitvec(const float* __restrict x, uint8_t* __restrict b, size_t d) {
 }
 
 /* Same but for n vectors.
-   Ensure that the ouptut b is byte-aligned (pad with 0s). */
+   Ensure that the output b is byte-aligned (pad with 0s). */
 void fvecs2bitvecs(
         const float* __restrict x,
         uint8_t* __restrict b,

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

@@ -166,9 +166,12 @@ struct HammingComputer20 {
     void set(const uint8_t* a8, int code_size) {
         assert(code_size == 20);
         const uint64_t* a = (uint64_t*)a8;
+        const uint32_t* b = (uint32_t*)a8;
         a0 = a[0];
         a1 = a[1];
-        a2 = a[2];
+        // can't read a[2] since it is uint64_t, not uint32_t
+        // results in AddressSanitizer failure reading past end of array
+        a2 = b[4];
     }
 
     inline int hamming(const uint8_t* b8) const {

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

@@ -55,7 +55,7 @@ SPECIALIZED_HC(64);
 /***************************************************************************
  * Dispatching function that takes a code size and a consumer object
  * the consumer object should contain a retun type t and a operation template
- * function f() that to be called to perform the operation.
+ * function f() that must be called to perform the operation.
  **************************************************************************/
 
 template <class Consumer, class... Types>
@@ -76,6 +76,7 @@ typename Consumer::T dispatch_HammingComputer(
         default:
             return consumer.template f<HammingComputerDefault>(args...);
     }
+#undef DISPATCH_HC
 }
 
 } // namespace faiss

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

@@ -54,6 +54,37 @@ double RandomGenerator::rand_double() {
     return mt() / double(mt.max());
 }
 
+SplitMix64RandomGenerator::SplitMix64RandomGenerator(int64_t seed)
+        : state{static_cast<uint64_t>(seed)} {}
+
+int SplitMix64RandomGenerator::rand_int() {
+    return next() & 0x7fffffff;
+}
+
+int64_t SplitMix64RandomGenerator::rand_int64() {
+    uint64_t value = next();
+    return static_cast<int64_t>(value & 0x7fffffffffffffffULL);
+}
+
+int SplitMix64RandomGenerator::rand_int(int max) {
+    return next() % max;
+}
+
+float SplitMix64RandomGenerator::rand_float() {
+    return next() / float(std::numeric_limits<uint64_t>::max());
+}
+
+double SplitMix64RandomGenerator::rand_double() {
+    return next() / double(std::numeric_limits<uint64_t>::max());
+}
+
+uint64_t SplitMix64RandomGenerator::next() {
+    uint64_t z = (state += 0x9e3779b97f4a7c15ULL);
+    z = (z ^ (z >> 30)) * 0xbf58476d1ce4e5b9ULL;
+    z = (z ^ (z >> 27)) * 0x94d049bb133111ebULL;
+    return z ^ (z >> 31);
+}
+
 /***********************************************************************
  * Random functions in this C file only exist because Torch
  *  counterparts are slow and not multi-threaded.  Typical use is for
@@ -162,6 +193,18 @@ void rand_perm(int* perm, size_t n, int64_t seed) {
     }
 }
 
+void rand_perm_splitmix64(int* perm, size_t n, int64_t seed) {
+    for (size_t i = 0; i < n; i++)
+        perm[i] = i;
+
+    SplitMix64RandomGenerator rng(seed);
+
+    for (size_t i = 0; i + 1 < n; i++) {
+        int i2 = i + rng.rand_int(n - i);
+        std::swap(perm[i], perm[i2]);
+    }
+}
+
 void byte_rand(uint8_t* x, size_t n, int64_t seed) {
     // only try to parallelize on large enough arrays
     const size_t nblock = n < 1024 ? 1 : 1024;

data/vendor/faiss/faiss/utils/random.h

@@ -43,6 +43,30 @@ struct RandomGenerator {
     explicit RandomGenerator(int64_t seed = 1234);
 };
 
+/// fast random generator that cannot be used in multithreaded contexts.
+/// based on https://prng.di.unimi.it/
+struct SplitMix64RandomGenerator {
+    uint64_t state;
+
+    /// random positive integer
+    int rand_int();
+
+    /// random int64_t
+    int64_t rand_int64();
+
+    /// generate random integer between 0 and max-1
+    int rand_int(int max);
+
+    /// between 0 and 1
+    float rand_float();
+
+    double rand_double();
+
+    explicit SplitMix64RandomGenerator(int64_t seed = 1234);
+
+    uint64_t next();
+};
+
 /* Generate an array of uniform random floats / multi-threaded implementation */
 void float_rand(float* x, size_t n, int64_t seed);
 void float_randn(float* x, size_t n, int64_t seed);
@@ -53,6 +77,7 @@ void int64_rand_max(int64_t* x, size_t n, uint64_t max, int64_t seed);
 
 /* random permutation */
 void rand_perm(int* perm, size_t n, int64_t seed);
+void rand_perm_splitmix64(int* perm, size_t n, int64_t seed);
 
 /* Random set of vectors with intrinsic dimensionality 10 that is harder to
  * index than a subspace of dim 10 but easier than uniform data in dimension d

data/vendor/faiss/faiss/utils/simdlib.h

@@ -14,7 +14,12 @@
  * functions.
  */
 
-#ifdef __AVX2__
+#if defined(__AVX512F__)
+
+#include <faiss/utils/simdlib_avx2.h>
+#include <faiss/utils/simdlib_avx512.h>
+
+#elif defined(__AVX2__)
 
 #include <faiss/utils/simdlib_avx2.h>
 
@@ -22,6 +27,10 @@
 
 #include <faiss/utils/simdlib_neon.h>
 
+#elif defined(__PPC64__)
+
+#include <faiss/utils/simdlib_ppc64.h>
+
 #else
 
 // emulated = all operations are implemented as scalars