faiss 0.3.0 → 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/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