faiss 0.3.0 → 0.3.2

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

@@ -5,13 +5,12 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-// -*- c++ -*-
-
 #include <faiss/utils/distances.h>
 
 #include <algorithm>
 #include <cassert>
 #include <cmath>
+#include <cstddef>
 #include <cstdio>
 #include <cstring>
 
@@ -64,7 +63,7 @@ void fvec_norms_L2(
         const float* __restrict x,
         size_t d,
         size_t nx) {
-#pragma omp parallel for
+#pragma omp parallel for if (nx > 10000)
     for (int64_t i = 0; i < nx; i++) {
         nr[i] = sqrtf(fvec_norm_L2sqr(x + i * d, d));
     }
@@ -75,24 +74,52 @@ void fvec_norms_L2sqr(
         const float* __restrict x,
         size_t d,
         size_t nx) {
-#pragma omp parallel for
+#pragma omp parallel for if (nx > 10000)
     for (int64_t i = 0; i < nx; i++)
         nr[i] = fvec_norm_L2sqr(x + i * d, d);
 }
 
-void fvec_renorm_L2(size_t d, size_t nx, float* __restrict x) {
-#pragma omp parallel for
+// The following is a workaround to a problem
+// in OpenMP in fbcode. The crash occurs
+// inside OMP when IndexIVFSpectralHash::set_query()
+// calls fvec_renorm_L2. set_query() is always
+// calling this function with nx == 1, so even
+// the omp version should run single threaded,
+// as per the if condition of the omp pragma.
+// Instead, the omp version crashes inside OMP.
+// The workaround below is explicitly branching
+// off to a codepath without omp.
+
+#define FVEC_RENORM_L2_IMPL                   \
+    float* __restrict xi = x + i * d;         \
+                                              \
+    float nr = fvec_norm_L2sqr(xi, d);        \
+                                              \
+    if (nr > 0) {                             \
+        size_t j;                             \
+        const float inv_nr = 1.0 / sqrtf(nr); \
+        for (j = 0; j < d; j++)               \
+            xi[j] *= inv_nr;                  \
+    }
+
+void fvec_renorm_L2_noomp(size_t d, size_t nx, float* __restrict x) {
     for (int64_t i = 0; i < nx; i++) {
-        float* __restrict xi = x + i * d;
+        FVEC_RENORM_L2_IMPL
+    }
+}
 
-        float nr = fvec_norm_L2sqr(xi, d);
+void fvec_renorm_L2_omp(size_t d, size_t nx, float* __restrict x) {
+#pragma omp parallel for if (nx > 10000)
+    for (int64_t i = 0; i < nx; i++) {
+        FVEC_RENORM_L2_IMPL
+    }
+}
 
-        if (nr > 0) {
-            size_t j;
-            const float inv_nr = 1.0 / sqrtf(nr);
-            for (j = 0; j < d; j++)
-                xi[j] *= inv_nr;
-        }
+void fvec_renorm_L2(size_t d, size_t nx, float* __restrict x) {
+    if (nx <= 10000) {
+        fvec_renorm_L2_noomp(d, nx, x);
+    } else {
+        fvec_renorm_L2_omp(d, nx, x);
     }
 }
 
@@ -103,19 +130,17 @@ void fvec_renorm_L2(size_t d, size_t nx, float* __restrict x) {
 namespace {
 
 /* Find the nearest neighbors for nx queries in a set of ny vectors */
-template <class ResultHandler, bool use_sel = false>
+template <class BlockResultHandler>
 void exhaustive_inner_product_seq(
         const float* x,
         const float* y,
         size_t d,
         size_t nx,
         size_t ny,
-        ResultHandler& res,
-        const IDSelector* sel = nullptr) {
-    using SingleResultHandler = typename ResultHandler::SingleResultHandler;
-    int nt = std::min(int(nx), omp_get_max_threads());
-
-    FAISS_ASSERT(use_sel == (sel != nullptr));
+        BlockResultHandler& res) {
+    using SingleResultHandler =
+            typename BlockResultHandler::SingleResultHandler;
+    [[maybe_unused]] int nt = std::min(int(nx), omp_get_max_threads());
 
 #pragma omp parallel num_threads(nt)
     {
@@ -128,7 +153,7 @@ void exhaustive_inner_product_seq(
             resi.begin(i);
 
             for (size_t j = 0; j < ny; j++, y_j += d) {
-                if (use_sel && !sel->is_member(j)) {
+                if (!res.is_in_selection(j)) {
                     continue;
                 }
                 float ip = fvec_inner_product(x_i, y_j, d);
@@ -139,19 +164,17 @@ void exhaustive_inner_product_seq(
     }
 }
 
-template <class ResultHandler, bool use_sel = false>
+template <class BlockResultHandler>
 void exhaustive_L2sqr_seq(
         const float* x,
         const float* y,
         size_t d,
         size_t nx,
         size_t ny,
-        ResultHandler& res,
-        const IDSelector* sel = nullptr) {
-    using SingleResultHandler = typename ResultHandler::SingleResultHandler;
-    int nt = std::min(int(nx), omp_get_max_threads());
-
-    FAISS_ASSERT(use_sel == (sel != nullptr));
+        BlockResultHandler& res) {
+    using SingleResultHandler =
+            typename BlockResultHandler::SingleResultHandler;
+    [[maybe_unused]] int nt = std::min(int(nx), omp_get_max_threads());
 
 #pragma omp parallel num_threads(nt)
     {
@@ -162,7 +185,7 @@ void exhaustive_L2sqr_seq(
             const float* y_j = y;
             resi.begin(i);
             for (size_t j = 0; j < ny; j++, y_j += d) {
-                if (use_sel && !sel->is_member(j)) {
+                if (!res.is_in_selection(j)) {
                     continue;
                 }
                 float disij = fvec_L2sqr(x_i, y_j, d);
@@ -174,14 +197,14 @@ void exhaustive_L2sqr_seq(
 }
 
 /** Find the nearest neighbors for nx queries in a set of ny vectors */
-template <class ResultHandler>
+template <class BlockResultHandler>
 void exhaustive_inner_product_blas(
         const float* x,
         const float* y,
         size_t d,
         size_t nx,
         size_t ny,
-        ResultHandler& res) {
+        BlockResultHandler& res) {
     // BLAS does not like empty matrices
     if (nx == 0 || ny == 0)
         return;
@@ -230,14 +253,14 @@ void exhaustive_inner_product_blas(
 
 // distance correction is an operator that can be applied to transform
 // the distances
-template <class ResultHandler>
+template <class BlockResultHandler>
 void exhaustive_L2sqr_blas_default_impl(
         const float* x,
         const float* y,
         size_t d,
         size_t nx,
         size_t ny,
-        ResultHandler& res,
+        BlockResultHandler& res,
         const float* y_norms = nullptr) {
     // BLAS does not like empty matrices
     if (nx == 0 || ny == 0)
@@ -297,6 +320,9 @@ void exhaustive_L2sqr_blas_default_impl(
                     float ip = *ip_line;
                     float dis = x_norms[i] + y_norms[j] - 2 * ip;
 
+                    if (!res.is_in_selection(j)) {
+                        dis = HUGE_VALF;
+                    }
                     // negative values can occur for identical vectors
                     // due to roundoff errors
                     if (dis < 0)
@@ -313,14 +339,14 @@ void exhaustive_L2sqr_blas_default_impl(
     }
 }
 
-template <class ResultHandler>
+template <class BlockResultHandler>
 void exhaustive_L2sqr_blas(
         const float* x,
         const float* y,
         size_t d,
         size_t nx,
         size_t ny,
-        ResultHandler& res,
+        BlockResultHandler& res,
         const float* y_norms = nullptr) {
     exhaustive_L2sqr_blas_default_impl(x, y, d, nx, ny, res);
 }
@@ -332,7 +358,7 @@ void exhaustive_L2sqr_blas_cmax_avx2(
         size_t d,
         size_t nx,
         size_t ny,
-        SingleBestResultHandler<CMax<float, int64_t>>& res,
+        Top1BlockResultHandler<CMax<float, int64_t>>& res,
         const float* y_norms) {
     // BLAS does not like empty matrices
     if (nx == 0 || ny == 0)
@@ -388,8 +414,8 @@ void exhaustive_L2sqr_blas_cmax_avx2(
             for (int64_t i = i0; i < i1; i++) {
                 float* ip_line = ip_block.get() + (i - i0) * (j1 - j0);
 
-                _mm_prefetch(ip_line, _MM_HINT_NTA);
-                _mm_prefetch(ip_line + 16, _MM_HINT_NTA);
+                _mm_prefetch((const char*)ip_line, _MM_HINT_NTA);
+                _mm_prefetch((const char*)(ip_line + 16), _MM_HINT_NTA);
 
                 // constant
                 const __m256 mul_minus2 = _mm256_set1_ps(-2);
@@ -416,8 +442,8 @@ void exhaustive_L2sqr_blas_cmax_avx2(
 
                 // process 16 elements per loop
                 for (; idx_j < (count / 16) * 16; idx_j += 16, ip_line += 16) {
-                    _mm_prefetch(ip_line + 32, _MM_HINT_NTA);
-                    _mm_prefetch(ip_line + 48, _MM_HINT_NTA);
+                    _mm_prefetch((const char*)(ip_line + 32), _MM_HINT_NTA);
+                    _mm_prefetch((const char*)(ip_line + 48), _MM_HINT_NTA);
 
                     // load values for norms
                     const __m256 y_norm_0 =
@@ -535,13 +561,13 @@ void exhaustive_L2sqr_blas_cmax_avx2(
 
 // an override if only a single closest point is needed
 template <>
-void exhaustive_L2sqr_blas<SingleBestResultHandler<CMax<float, int64_t>>>(
+void exhaustive_L2sqr_blas<Top1BlockResultHandler<CMax<float, int64_t>>>(
         const float* x,
         const float* y,
         size_t d,
         size_t nx,
         size_t ny,
-        SingleBestResultHandler<CMax<float, int64_t>>& res,
+        Top1BlockResultHandler<CMax<float, int64_t>>& res,
         const float* y_norms) {
 #if defined(__AVX2__)
     // use a faster fused kernel if available
@@ -562,34 +588,50 @@ void exhaustive_L2sqr_blas<SingleBestResultHandler<CMax<float, int64_t>>>(
 
     // run the default implementation
     exhaustive_L2sqr_blas_default_impl<
-            SingleBestResultHandler<CMax<float, int64_t>>>(
+            Top1BlockResultHandler<CMax<float, int64_t>>>(
             x, y, d, nx, ny, res, y_norms);
 #else
     // run the default implementation
     exhaustive_L2sqr_blas_default_impl<
-            SingleBestResultHandler<CMax<float, int64_t>>>(
+            Top1BlockResultHandler<CMax<float, int64_t>>>(
             x, y, d, nx, ny, res, y_norms);
 #endif
 }
 
-template <class ResultHandler>
-void knn_L2sqr_select(
-        const float* x,
-        const float* y,
-        size_t d,
-        size_t nx,
-        size_t ny,
-        ResultHandler& res,
-        const float* y_norm2,
-        const IDSelector* sel) {
-    if (sel) {
-        exhaustive_L2sqr_seq<ResultHandler, true>(x, y, d, nx, ny, res, sel);
-    } else if (nx < distance_compute_blas_threshold) {
-        exhaustive_L2sqr_seq(x, y, d, nx, ny, res);
-    } else {
-        exhaustive_L2sqr_blas(x, y, d, nx, ny, res, y_norm2);
+struct Run_search_inner_product {
+    using T = void;
+    template <class BlockResultHandler>
+    void f(BlockResultHandler& res,
+           const float* x,
+           const float* y,
+           size_t d,
+           size_t nx,
+           size_t ny) {
+        if (res.sel || nx < distance_compute_blas_threshold) {
+            exhaustive_inner_product_seq(x, y, d, nx, ny, res);
+        } else {
+            exhaustive_inner_product_blas(x, y, d, nx, ny, res);
+        }
     }
-}
+};
+
+struct Run_search_L2sqr {
+    using T = void;
+    template <class BlockResultHandler>
+    void f(BlockResultHandler& res,
+           const float* x,
+           const float* y,
+           size_t d,
+           size_t nx,
+           size_t ny,
+           const float* y_norm2) {
+        if (res.sel || nx < distance_compute_blas_threshold) {
+            exhaustive_L2sqr_seq(x, y, d, nx, ny, res);
+        } else {
+            exhaustive_L2sqr_blas(x, y, d, nx, ny, res, y_norm2);
+        }
+    }
+};
 
 } // anonymous namespace
 
@@ -609,7 +651,7 @@ void knn_inner_product(
         size_t nx,
         size_t ny,
         size_t k,
-        float* val,
+        float* vals,
         int64_t* ids,
         const IDSelector* sel) {
     int64_t imin = 0;
@@ -622,30 +664,14 @@ void knn_inner_product(
     }
     if (auto sela = dynamic_cast<const IDSelectorArray*>(sel)) {
         knn_inner_products_by_idx(
-                x, y, sela->ids, d, nx, sela->n, k, val, ids, 0);
+                x, y, sela->ids, d, nx, ny, sela->n, k, vals, ids, 0);
         return;
     }
-    if (k < distance_compute_min_k_reservoir) {
-        using RH = HeapResultHandler<CMin<float, int64_t>>;
-        RH res(nx, val, ids, k);
-        if (sel) {
-            exhaustive_inner_product_seq<RH, true>(x, y, d, nx, ny, res, sel);
-        } else if (nx < distance_compute_blas_threshold) {
-            exhaustive_inner_product_seq(x, y, d, nx, ny, res);
-        } else {
-            exhaustive_inner_product_blas(x, y, d, nx, ny, res);
-        }
-    } else {
-        using RH = ReservoirResultHandler<CMin<float, int64_t>>;
-        RH res(nx, val, ids, k);
-        if (sel) {
-            exhaustive_inner_product_seq<RH, true>(x, y, d, nx, ny, res, sel);
-        } else if (nx < distance_compute_blas_threshold) {
-            exhaustive_inner_product_seq(x, y, d, nx, ny, res, nullptr);
-        } else {
-            exhaustive_inner_product_blas(x, y, d, nx, ny, res);
-        }
-    }
+
+    Run_search_inner_product r;
+    dispatch_knn_ResultHandler(
+            nx, vals, ids, k, METRIC_INNER_PRODUCT, sel, r, x, y, d, nx, ny);
+
     if (imin != 0) {
         for (size_t i = 0; i < nx * k; i++) {
             if (ids[i] >= 0) {
@@ -687,19 +713,14 @@ void knn_L2sqr(
         sel = nullptr;
     }
     if (auto sela = dynamic_cast<const IDSelectorArray*>(sel)) {
-        knn_L2sqr_by_idx(x, y, sela->ids, d, nx, sela->n, k, vals, ids, 0);
+        knn_L2sqr_by_idx(x, y, sela->ids, d, nx, ny, sela->n, k, vals, ids, 0);
         return;
     }
-    if (k == 1) {
-        SingleBestResultHandler<CMax<float, int64_t>> res(nx, vals, ids);
-        knn_L2sqr_select(x, y, d, nx, ny, res, y_norm2, sel);
-    } else if (k < distance_compute_min_k_reservoir) {
-        HeapResultHandler<CMax<float, int64_t>> res(nx, vals, ids, k);
-        knn_L2sqr_select(x, y, d, nx, ny, res, y_norm2, sel);
-    } else {
-        ReservoirResultHandler<CMax<float, int64_t>> res(nx, vals, ids, k);
-        knn_L2sqr_select(x, y, d, nx, ny, res, y_norm2, sel);
-    }
+
+    Run_search_L2sqr r;
+    dispatch_knn_ResultHandler(
+            nx, vals, ids, k, METRIC_L2, sel, r, x, y, d, nx, ny, y_norm2);
+
     if (imin != 0) {
         for (size_t i = 0; i < nx * k; i++) {
             if (ids[i] >= 0) {
@@ -726,6 +747,7 @@ void knn_L2sqr(
  * Range search
  ***************************************************************************/
 
+// TODO accept a y_norm2 as well
 void range_search_L2sqr(
         const float* x,
         const float* y,
@@ -735,15 +757,9 @@ void range_search_L2sqr(
         float radius,
         RangeSearchResult* res,
         const IDSelector* sel) {
-    using RH = RangeSearchResultHandler<CMax<float, int64_t>>;
-    RH resh(res, radius);
-    if (sel) {
-        exhaustive_L2sqr_seq<RH, true>(x, y, d, nx, ny, resh, sel);
-    } else if (nx < distance_compute_blas_threshold) {
-        exhaustive_L2sqr_seq(x, y, d, nx, ny, resh, sel);
-    } else {
-        exhaustive_L2sqr_blas(x, y, d, nx, ny, resh);
-    }
+    Run_search_L2sqr r;
+    dispatch_range_ResultHandler(
+            res, radius, METRIC_L2, sel, r, x, y, d, nx, ny, nullptr);
 }
 
 void range_search_inner_product(
@@ -755,15 +771,9 @@ void range_search_inner_product(
         float radius,
         RangeSearchResult* res,
         const IDSelector* sel) {
-    using RH = RangeSearchResultHandler<CMin<float, int64_t>>;
-    RH resh(res, radius);
-    if (sel) {
-        exhaustive_inner_product_seq<RH, true>(x, y, d, nx, ny, resh, sel);
-    } else if (nx < distance_compute_blas_threshold) {
-        exhaustive_inner_product_seq(x, y, d, nx, ny, resh);
-    } else {
-        exhaustive_inner_product_blas(x, y, d, nx, ny, resh);
-    }
+    Run_search_inner_product r;
+    dispatch_range_ResultHandler(
+            res, radius, METRIC_INNER_PRODUCT, sel, r, x, y, d, nx, ny);
 }
 
 /***************************************************************************
@@ -786,9 +796,11 @@ void fvec_inner_products_by_idx(
         const float* xj = x + j * d;
         float* __restrict ipj = ip + j * ny;
         for (size_t i = 0; i < ny; i++) {
-            if (idsj[i] < 0)
-                continue;
-            ipj[i] = fvec_inner_product(xj, y + d * idsj[i], d);
+            if (idsj[i] < 0) {
+                ipj[i] = -INFINITY;
+            } else {
+                ipj[i] = fvec_inner_product(xj, y + d * idsj[i], d);
+            }
         }
     }
 }
@@ -809,9 +821,11 @@ void fvec_L2sqr_by_idx(
         const float* xj = x + j * d;
         float* __restrict disj = dis + j * ny;
         for (size_t i = 0; i < ny; i++) {
-            if (idsj[i] < 0)
-                continue;
-            disj[i] = fvec_L2sqr(xj, y + d * idsj[i], d);
+            if (idsj[i] < 0) {
+                disj[i] = INFINITY;
+            } else {
+                disj[i] = fvec_L2sqr(xj, y + d * idsj[i], d);
+            }
         }
     }
 }
@@ -828,6 +842,8 @@ void pairwise_indexed_L2sqr(
     for (int64_t j = 0; j < n; j++) {
         if (ix[j] >= 0 && iy[j] >= 0) {
             dis[j] = fvec_L2sqr(x + d * ix[j], y + d * iy[j], d);
+        } else {
+            dis[j] = INFINITY;
         }
     }
 }
@@ -844,6 +860,8 @@ void pairwise_indexed_inner_product(
     for (int64_t j = 0; j < n; j++) {
         if (ix[j] >= 0 && iy[j] >= 0) {
             dis[j] = fvec_inner_product(x + d * ix[j], y + d * iy[j], d);
+        } else {
+            dis[j] = -INFINITY;
         }
     }
 }
@@ -857,6 +875,7 @@ void knn_inner_products_by_idx(
         size_t d,
         size_t nx,
         size_t ny,
+        size_t nsubset,
         size_t k,
         float* res_vals,
         int64_t* res_ids,
@@ -874,9 +893,10 @@ void knn_inner_products_by_idx(
         int64_t* __restrict idxi = res_ids + i * k;
         minheap_heapify(k, simi, idxi);
 
-        for (j = 0; j < ny; j++) {
-            if (idsi[j] < 0)
+        for (j = 0; j < nsubset; j++) {
+            if (idsi[j] < 0 || idsi[j] >= ny) {
                 break;
+            }
             float ip = fvec_inner_product(x_, y + d * idsi[j], d);
 
             if (ip > simi[0]) {
@@ -894,6 +914,7 @@ void knn_L2sqr_by_idx(
         size_t d,
         size_t nx,
         size_t ny,
+        size_t nsubset,
         size_t k,
         float* res_vals,
         int64_t* res_ids,
@@ -908,7 +929,10 @@ void knn_L2sqr_by_idx(
         float* __restrict simi = res_vals + i * k;
         int64_t* __restrict idxi = res_ids + i * k;
         maxheap_heapify(k, simi, idxi);
-        for (size_t j = 0; j < ny; j++) {
+        for (size_t j = 0; j < nsubset; j++) {
+            if (idsi[j] < 0 || idsi[j] >= ny) {
+                break;
+            }
             float disij = fvec_L2sqr(x_, y + d * idsi[j], d);
 
             if (disij < simi[0]) {