faiss 0.5.2 → 0.6.0

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

@@ -12,217 +12,224 @@
 
 #include <faiss/MetricType.h>
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/simd_dispatch.h>
 #include <faiss/utils/distances.h>
 #include <cmath>
 #include <type_traits>
 
 namespace faiss {
 
+/***************************************************************************
+ * VectorDistance base class - contains common data members and type defs
+ **************************************************************************/
+
 template <MetricType mt>
-struct VectorDistance {
+struct VectorDistanceBase {
     size_t d;
     float metric_arg;
+    static constexpr MetricType metric = mt;
     static constexpr bool is_similarity = is_similarity_metric(mt);
 
-    inline float operator()(const float* x, const float* y) const;
-
-    // heap template to use for this type of metric
     using C = typename std::conditional<
             is_similarity_metric(mt),
             CMin<float, int64_t>,
             CMax<float, int64_t>>::type;
 };
 
-template <>
-inline float VectorDistance<METRIC_L2>::operator()(
-        const float* x,
-        const float* y) const {
-    return fvec_L2sqr(x, y, d);
-}
+/***************************************************************************
+ * VectorDistance struct template - specializations for each metric type
+ **************************************************************************/
 
-template <>
-inline float VectorDistance<METRIC_INNER_PRODUCT>::operator()(
-        const float* x,
-        const float* y) const {
-    return fvec_inner_product(x, y, d);
-}
+template <MetricType mt, SIMDLevel level>
+struct VectorDistance : VectorDistanceBase<mt> {
+    inline float operator()(const float* x, const float* y) const;
+};
 
-template <>
-inline float VectorDistance<METRIC_L1>::operator()(
-        const float* x,
-        const float* y) const {
-    return fvec_L1(x, y, d);
-}
+template <SIMDLevel level>
+struct VectorDistance<METRIC_L2, level> : VectorDistanceBase<METRIC_L2> {
+    inline float operator()(const float* x, const float* y) const {
+        return fvec_L2sqr<level>(x, y, this->d);
+    }
+};
 
-template <>
-inline float VectorDistance<METRIC_Linf>::operator()(
-        const float* x,
-        const float* y) const {
-    return fvec_Linf(x, y, d);
-    /*
-        float vmax = 0;
-        for (size_t i = 0; i < d; i++) {
-            float diff = fabs (x[i] - y[i]);
-            if (diff > vmax) vmax = diff;
-        }
-     return vmax;*/
-}
+template <SIMDLevel level>
+struct VectorDistance<METRIC_INNER_PRODUCT, level>
+        : VectorDistanceBase<METRIC_INNER_PRODUCT> {
+    inline float operator()(const float* x, const float* y) const {
+        return fvec_inner_product<level>(x, y, this->d);
+    }
+};
 
-template <>
-inline float VectorDistance<METRIC_Lp>::operator()(
-        const float* x,
-        const float* y) const {
-    float accu = 0;
-    for (size_t i = 0; i < d; i++) {
-        float diff = fabs(x[i] - y[i]);
-        accu += powf(diff, metric_arg);
+template <SIMDLevel level>
+struct VectorDistance<METRIC_L1, level> : VectorDistanceBase<METRIC_L1> {
+    inline float operator()(const float* x, const float* y) const {
+        return fvec_L1<level>(x, y, this->d);
     }
-    return accu;
-}
+};
 
-template <>
-inline float VectorDistance<METRIC_Canberra>::operator()(
-        const float* x,
-        const float* y) const {
-    float accu = 0;
-    for (size_t i = 0; i < d; i++) {
-        float xi = x[i], yi = y[i];
-        accu += fabs(xi - yi) / (fabs(xi) + fabs(yi));
+template <SIMDLevel level>
+struct VectorDistance<METRIC_Linf, level> : VectorDistanceBase<METRIC_Linf> {
+    inline float operator()(const float* x, const float* y) const {
+        return fvec_Linf<level>(x, y, this->d);
     }
-    return accu;
-}
+};
 
 template <>
-inline float VectorDistance<METRIC_BrayCurtis>::operator()(
-        const float* x,
-        const float* y) const {
-    float accu_num = 0, accu_den = 0;
-    for (size_t i = 0; i < d; i++) {
-        float xi = x[i], yi = y[i];
-        accu_num += fabs(xi - yi);
-        accu_den += fabs(xi + yi);
+struct VectorDistance<METRIC_Lp, SIMDLevel::NONE>
+        : VectorDistanceBase<METRIC_Lp> {
+    inline float operator()(const float* x, const float* y) const {
+        float accu = 0;
+        for (size_t i = 0; i < this->d; i++) {
+            float diff = fabs(x[i] - y[i]);
+            accu += powf(diff, this->metric_arg);
+        }
+        return accu;
     }
-    return accu_num / accu_den;
-}
+};
 
 template <>
-inline float VectorDistance<METRIC_JensenShannon>::operator()(
-        const float* x,
-        const float* y) const {
-    float accu = 0;
-    for (size_t i = 0; i < d; i++) {
-        float xi = x[i], yi = y[i];
-        float mi = 0.5 * (xi + yi);
-        float kl1 = -xi * log(mi / xi);
-        float kl2 = -yi * log(mi / yi);
-        accu += kl1 + kl2;
+struct VectorDistance<METRIC_Canberra, SIMDLevel::NONE>
+        : VectorDistanceBase<METRIC_Canberra> {
+    inline float operator()(const float* x, const float* y) const {
+        float accu = 0;
+        for (size_t i = 0; i < this->d; i++) {
+            float xi = x[i], yi = y[i];
+            accu += fabs(xi - yi) / (fabs(xi) + fabs(yi));
+        }
+        return accu;
     }
-    return 0.5 * accu;
-}
+};
 
 template <>
-inline float VectorDistance<METRIC_Jaccard>::operator()(
-        const float* x,
-        const float* y) const {
-    // WARNING: this distance is defined only for positive input vectors.
-    // Providing vectors with negative values would lead to incorrect results.
-    float accu_num = 0, accu_den = 0;
-    for (size_t i = 0; i < d; i++) {
-        accu_num += fmin(x[i], y[i]);
-        accu_den += fmax(x[i], y[i]);
+struct VectorDistance<METRIC_BrayCurtis, SIMDLevel::NONE>
+        : VectorDistanceBase<METRIC_BrayCurtis> {
+    inline float operator()(const float* x, const float* y) const {
+        float accu_num = 0, accu_den = 0;
+        for (size_t i = 0; i < this->d; i++) {
+            float xi = x[i], yi = y[i];
+            accu_num += fabs(xi - yi);
+            accu_den += fabs(xi + yi);
+        }
+        return accu_num / accu_den;
     }
-    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++;
+struct VectorDistance<METRIC_JensenShannon, SIMDLevel::NONE>
+        : VectorDistanceBase<METRIC_JensenShannon> {
+    inline float operator()(const float* x, const float* y) const {
+        float accu = 0;
+        for (size_t i = 0; i < this->d; i++) {
+            float xi = x[i], yi = y[i];
+            float mi = 0.5 * (xi + yi);
+            float kl1 = -xi * log(mi / xi);
+            float kl2 = -yi * log(mi / yi);
+            accu += kl1 + kl2;
         }
+        return 0.5 * accu;
     }
-    if (present == 0) {
-        return NAN;
+};
+
+template <>
+struct VectorDistance<METRIC_Jaccard, SIMDLevel::NONE>
+        : VectorDistanceBase<METRIC_Jaccard> {
+    inline float operator()(const float* x, const float* y) const {
+        // WARNING: this distance is defined only for positive input vectors.
+        // Providing vectors with negative values would lead to incorrect
+        // results.
+        float accu_num = 0, accu_den = 0;
+        for (size_t i = 0; i < this->d; i++) {
+            accu_num += fmin(x[i], y[i]);
+            accu_den += fmax(x[i], y[i]);
+        }
+        return accu_num / accu_den;
     }
-    return float(d) / float(present) * accu;
-}
+};
 
 template <>
-inline float VectorDistance<METRIC_GOWER>::operator()(
-        const float* x,
-        const float* y) const {
-    float accu = 0;
-    size_t valid_dims = 0;
-
-    for (size_t i = 0; i < d; i++) {
-        if (std::isnan(x[i]) || std::isnan(y[i])) {
-            continue;
+struct VectorDistance<METRIC_NaNEuclidean, SIMDLevel::NONE>
+        : VectorDistanceBase<METRIC_NaNEuclidean> {
+    inline float 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 < this->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(this->d) / float(present) * accu;
+    }
+};
 
-        if (x[i] >= 0 && y[i] >= 0) {
-            if (x[i] > 1 || y[i] > 1) {
+template <>
+struct VectorDistance<METRIC_GOWER, SIMDLevel::NONE>
+        : VectorDistanceBase<METRIC_GOWER> {
+    inline float operator()(const float* x, const float* y) const {
+        float accu = 0;
+        size_t valid_dims = 0;
+
+        for (size_t i = 0; i < this->d; i++) {
+            if (std::isnan(x[i]) || std::isnan(y[i])) {
+                continue;
+            }
+
+            if (x[i] >= 0 && y[i] >= 0) {
+                if (x[i] > 1 || y[i] > 1) {
+                    return std::numeric_limits<float>::quiet_NaN();
+                }
+                accu += fabs(x[i] - y[i]);
+            } else if (x[i] < 0 && y[i] < 0) {
+                accu += float(int(x[i] != y[i]));
+            } else {
                 return std::numeric_limits<float>::quiet_NaN();
             }
-            // Numeric dimensions are in [0,1]
-            accu += fabs(x[i] - y[i]);
-        } else if (x[i] < 0 && y[i] < 0) {
-            // Categorical dimensions are negative values
-            accu += float(int(x[i] != y[i]));
-        } else {
-            // Invalid representation
-            return std::numeric_limits<float>::quiet_NaN();
+            valid_dims++;
         }
-        valid_dims++;
-    }
 
-    if (valid_dims == 0) {
-        return std::numeric_limits<float>::quiet_NaN();
+        if (valid_dims == 0) {
+            return std::numeric_limits<float>::quiet_NaN();
+        }
+        return accu / valid_dims;
     }
-    return accu / valid_dims;
-}
+};
 
 /***************************************************************************
- * Dispatching function that takes a metric type and a consumer object
- * the consumer object should contain a return 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.
+ * Dispatching function that takes a lambda directly.
+ * The lambda should be templated on VectorDistance, eg.:
+ *
+ *   auto result = with_VectorDistance(
+ *       metric, metric_arg, [&]<class VD>(VD vd) {
+ *           return vd(x, y);
+ *       });
  **************************************************************************/
 
-template <class Consumer, class... Types>
-typename Consumer::T dispatch_VectorDistance(
+template <typename LambdaType>
+auto with_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_GOWER);
-        default:
-            FAISS_THROW_FMT("Invalid metric %d", metric);
-    }
-#undef DISPATCH_VD
+        LambdaType&& action) {
+    auto dispatch_metric = [&]<MetricType mt>() {
+        auto call = [&]<SIMDLevel level>() {
+            VectorDistance<mt, level> vd = {d, metric_arg};
+            return action(vd);
+        };
+
+        constexpr bool has_simd = mt == METRIC_INNER_PRODUCT ||
+                mt == METRIC_L2 || mt == METRIC_L1 || mt == METRIC_Linf;
+        if constexpr (!has_simd) {
+            return call.template operator()<SIMDLevel::NONE>();
+        } else {
+            DISPATCH_SIMDLevel(call.template operator());
+        }
+    };
+    return with_metric_type(metric, dispatch_metric);
 }
 
 } // namespace faiss

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

@@ -11,10 +11,10 @@
 
 #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 +25,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;
@@ -128,19 +56,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 +83,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 +108,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 < 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 +150,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

@@ -13,6 +13,7 @@
 #include <stdint.h>
 
 #include <faiss/Index.h>
+#include <faiss/impl/IDSelector.h>
 
 #include <faiss/utils/Heap.h>
 
@@ -43,7 +44,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 +56,50 @@ FlatCodesDistanceComputer* get_extra_distance_computer(
         size_t nb,
         const float* xb);
 
+/// 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
+
 } // namespace faiss
 
 #include <faiss/utils/extra_distances-inl.h>

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

@@ -312,7 +312,6 @@ struct HammingComputerDefault {
             const uint8_t* a = a8 + 8 * quotient8;
             const uint8_t* b = b8 + 8 * quotient8;
             switch (remainder8) {
-                [[fallthrough]];
                 case 7:
                     accu += hamdis_tab_ham_bytes[a[6] ^ b[6]];
                     [[fallthrough]];

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

@@ -627,7 +627,7 @@ uint16_t simd_partition_fuzzy_with_bounds_histogram(
             n_lt = sum_below - hist[i];
             n_gt = n - sum_below;
         } else {
-            assert(!"not implemented");
+            assert(false && "not implemented");
         }
 
         IFV printf(