faiss 0.5.3 → 0.6.1

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

@@ -7,14 +7,15 @@
 
 // -*- c++ -*-
 
+#include <faiss/utils/distances_dispatch.h>
 #include <faiss/utils/extra_distances.h>
 
 #include <omp.h>
 #include <algorithm>
-#include <cmath>
 
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/DistanceComputer.h>
+#include <faiss/impl/IDSelector.h>
 #include <faiss/utils/utils.h>
 
 namespace faiss {
@@ -25,78 +26,6 @@ namespace faiss {
 
 namespace {
 
-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;
-            }
-        }
-    }
-};
-
-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 = 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;
-                }
-                // maxheap_reorder(k, simi, idxi);
-                heap_reorder<C>(k, simi, idxi);
-            }
-            InterruptCallback::check();
-        }
-    }
-};
-
 template <class VD>
 struct ExtraDistanceComputer : FlatCodesDistanceComputer {
     VD vd;
@@ -113,14 +42,14 @@ struct ExtraDistanceComputer : FlatCodesDistanceComputer {
     }
 
     ExtraDistanceComputer(
-            const VD& vd,
+            const VD& vd_in,
             const float* xb,
-            size_t nb,
-            const float* q = nullptr)
-            : FlatCodesDistanceComputer((uint8_t*)xb, vd.d * sizeof(float)),
-              vd(vd),
-              nb(nb),
-              q(q),
+            size_t nb_in,
+            const float* q_in = nullptr)
+            : FlatCodesDistanceComputer((uint8_t*)xb, vd_in.d * sizeof(float)),
+              vd(vd_in),
+              nb(nb_in),
+              q(q_in),
               b(xb) {}
 
     void set_query(const float* x) override {
@@ -128,19 +57,6 @@ 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(
@@ -168,9 +84,19 @@ void pairwise_extra_distances(
         ldd = nb;
     }
 
-    Run_pairwise_extra_distances run;
-    dispatch_VectorDistance(
-            d, mt, metric_arg, run, nq, xq, nb, xb, dis, ldq, ldb, ldd);
+    with_VectorDistance(d, mt, metric_arg, [&](auto vd) {
+#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;
+            }
+        }
+    });
 }
 
 void knn_extra_metrics(
@@ -183,10 +109,40 @@ void knn_extra_metrics(
         float metric_arg,
         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);
+        int64_t* indexes,
+        const IDSelector* sel) {
+    with_VectorDistance(d, mt, metric_arg, [&](auto vd) {
+        using C = typename decltype(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 < static_cast<int64_t>(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 = indexes + k * i;
+
+                heap_heapify<C>(k, simi, idxi);
+                for (j = 0; j < ny; j++) {
+                    if (!sel || sel->is_member(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;
+                }
+                heap_reorder<C>(k, simi, idxi);
+            }
+            InterruptCallback::check();
+        }
+    });
 }
 
 FlatCodesDistanceComputer* get_extra_distance_computer(
@@ -195,8 +151,10 @@ FlatCodesDistanceComputer* get_extra_distance_computer(
         float metric_arg,
         size_t nb,
         const float* xb) {
-    Run_get_distance_computer run;
-    return dispatch_VectorDistance(d, mt, metric_arg, run, xb, nb);
+    return with_VectorDistance(
+            d, mt, metric_arg, [&](auto vd) -> FlatCodesDistanceComputer* {
+                return new ExtraDistanceComputer<decltype(vd)>(vd, xb, nb);
+            });
 }
 
 } // namespace faiss

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

@@ -10,11 +10,13 @@
 /** In this file are the implementations of extra metrics beyond L2
  *  and inner product */
 
-#include <stdint.h>
+#include <cstdint>
 
-#include <faiss/Index.h>
-
-#include <faiss/utils/Heap.h>
+#include <faiss/MetricType.h>
+#include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/IDSelector.h>
+#include <faiss/utils/ordered_key_value.h>
+#include <faiss/utils/simd_levels.h>
 
 namespace faiss {
 
@@ -43,7 +45,8 @@ void knn_extra_metrics(
         float metric_arg,
         size_t k,
         float* distances,
-        int64_t* indexes);
+        int64_t* indexes,
+        const IDSelector* sel = nullptr);
 
 /** get a DistanceComputer that refers to this type of distance and
  *  indexes a flat array of size nb */
@@ -54,6 +57,75 @@ FlatCodesDistanceComputer* get_extra_distance_computer(
         size_t nb,
         const float* xb);
 
-} // namespace faiss
+/// Dispatch to a lambda with MetricType as a compile-time constant.
+/// This allows writing generic code that works with different metrics
+/// while maintaining compile-time optimization.
+///
+/// Example usage:
+///   auto result = with_metric_type(runtime_metric, [&](auto metric_tag) {
+///       constexpr MetricType M = decltype(metric_tag)::value;
+///       return compute_distance<M>(x, y);
+///   });
+#ifndef SWIG
+
+template <typename LambdaType>
+inline auto with_metric_type(MetricType metric, LambdaType&& action) {
+    switch (metric) {
+        case METRIC_INNER_PRODUCT:
+            return action.template operator()<METRIC_INNER_PRODUCT>();
+        case METRIC_L2:
+            return action.template operator()<METRIC_L2>();
+        case METRIC_L1:
+            return action.template operator()<METRIC_L1>();
+        case METRIC_Linf:
+            return action.template operator()<METRIC_Linf>();
+        case METRIC_Lp:
+            return action.template operator()<METRIC_Lp>();
+        case METRIC_Canberra:
+            return action.template operator()<METRIC_Canberra>();
+        case METRIC_BrayCurtis:
+            return action.template operator()<METRIC_BrayCurtis>();
+        case METRIC_JensenShannon:
+            return action.template operator()<METRIC_JensenShannon>();
+        case METRIC_Jaccard:
+            return action.template operator()<METRIC_Jaccard>();
+        case METRIC_NaNEuclidean:
+            return action.template operator()<METRIC_NaNEuclidean>();
+        case METRIC_GOWER:
+            return action.template operator()<METRIC_GOWER>();
+        default:
+            FAISS_THROW_FMT(
+                    "with_metric_type called with unknown metric %d",
+                    int(metric));
+    }
+}
+#endif // SWIG
+
+#ifndef SWIG
+
+/***************************************************************************
+ * VectorDistance base class - contains common data members and type defs
+ * VectorDistance struct template - specializations for each metric type
+ **************************************************************************/
+
+template <MetricType mt, SIMDLevel level>
+struct VectorDistance {
+    size_t d;
+    float metric_arg;
 
-#include <faiss/utils/extra_distances-inl.h>
+    VectorDistance(size_t d, float metric_arg) : d(d), metric_arg(metric_arg) {}
+
+    static constexpr MetricType metric = mt;
+    static constexpr bool is_similarity = is_similarity_metric(mt);
+
+    using C = typename std::conditional<
+            is_similarity_metric(mt),
+            CMin<float, int64_t>,
+            CMax<float, int64_t>>::type;
+
+    float operator()(const float* x, const float* y) const;
+};
+
+#endif // SWIG
+
+} // namespace faiss

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

@@ -10,9 +10,9 @@
 namespace faiss {
 
 // BitstringWriter and BitstringReader functions
-inline BitstringWriter::BitstringWriter(uint8_t* code, size_t code_size)
-        : code(code), code_size(code_size), i(0) {
-    memset(code, 0, code_size);
+inline BitstringWriter::BitstringWriter(uint8_t* code_in, size_t code_size_in)
+        : code(code_in), code_size(code_size_in), i(0) {
+    memset(code_in, 0, code_size_in);
 }
 
 inline void BitstringWriter::write(uint64_t x, int nbit) {
@@ -36,8 +36,10 @@ inline void BitstringWriter::write(uint64_t x, int nbit) {
     }
 }
 
-inline BitstringReader::BitstringReader(const uint8_t* code, size_t code_size)
-        : code(code), code_size(code_size), i(0) {}
+inline BitstringReader::BitstringReader(
+        const uint8_t* code_in,
+        size_t code_size_in)
+        : code(code_in), code_size(code_size_in), i(0) {}
 
 inline uint64_t BitstringReader::read(int nbit) {
     assert(code_size * 8 >= nbit + i);
@@ -85,18 +87,18 @@ struct HCounterState {
     int k;
 
     HCounterState(
-            int* counters,
-            int64_t* ids_per_dis,
+            int* counters_in,
+            int64_t* ids_per_dis_in,
             const uint8_t* x,
             int d,
-            int k)
-            : counters(counters),
-              ids_per_dis(ids_per_dis),
+            int k_in)
+            : counters(counters_in),
+              ids_per_dis(ids_per_dis_in),
               hc(x, d / 8),
               thres(d + 1),
               count_lt(0),
               count_eq(0),
-              k(k) {}
+              k(k_in) {}
 
     void update_counter(const uint8_t* y, size_t j) {
         int32_t dis = hc.hamming(y);