faiss 0.2.4 → 0.2.5

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

@@ -17,8 +17,13 @@
 
 #include <omp.h>
 
+#ifdef __AVX2__
+#include <immintrin.h>
+#endif
+
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/IDSelector.h>
 #include <faiss/impl/ResultHandler.h>
 
 #ifndef FINTEGER
@@ -96,17 +101,20 @@ 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>
+template <class ResultHandler, bool use_sel = false>
 void exhaustive_inner_product_seq(
         const float* x,
         const float* y,
         size_t d,
         size_t nx,
         size_t ny,
-        ResultHandler& res) {
+        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));
+
 #pragma omp parallel num_threads(nt)
     {
         SingleResultHandler resi(res);
@@ -117,27 +125,32 @@ void exhaustive_inner_product_seq(
 
             resi.begin(i);
 
-            for (size_t j = 0; j < ny; j++) {
+            for (size_t j = 0; j < ny; j++, y_j += d) {
+                if (use_sel && !sel->is_member(j)) {
+                    continue;
+                }
                 float ip = fvec_inner_product(x_i, y_j, d);
                 resi.add_result(ip, j);
-                y_j += d;
             }
             resi.end();
         }
     }
 }
 
-template <class ResultHandler>
+template <class ResultHandler, bool use_sel = false>
 void exhaustive_L2sqr_seq(
         const float* x,
         const float* y,
         size_t d,
         size_t nx,
         size_t ny,
-        ResultHandler& res) {
+        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));
+
 #pragma omp parallel num_threads(nt)
     {
         SingleResultHandler resi(res);
@@ -146,10 +159,12 @@ void exhaustive_L2sqr_seq(
             const float* x_i = x + i * d;
             const float* y_j = y;
             resi.begin(i);
-            for (size_t j = 0; j < ny; j++) {
+            for (size_t j = 0; j < ny; j++, y_j += d) {
+                if (use_sel && !sel->is_member(j)) {
+                    continue;
+                }
                 float disij = fvec_L2sqr(x_i, y_j, d);
                 resi.add_result(disij, j);
-                y_j += d;
             }
             resi.end();
         }
@@ -296,6 +311,232 @@ void exhaustive_L2sqr_blas(
     }
 }
 
+#ifdef __AVX2__
+// an override for AVX2 if only a single closest point is needed.
+template <>
+void exhaustive_L2sqr_blas<SingleBestResultHandler<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,
+        const float* y_norms) {
+    // BLAS does not like empty matrices
+    if (nx == 0 || ny == 0)
+        return;
+
+    /* block sizes */
+    const size_t bs_x = distance_compute_blas_query_bs;
+    const size_t bs_y = distance_compute_blas_database_bs;
+    // const size_t bs_x = 16, bs_y = 16;
+    std::unique_ptr<float[]> ip_block(new float[bs_x * bs_y]);
+    std::unique_ptr<float[]> x_norms(new float[nx]);
+    std::unique_ptr<float[]> del2;
+
+    fvec_norms_L2sqr(x_norms.get(), x, d, nx);
+
+    if (!y_norms) {
+        float* y_norms2 = new float[ny];
+        del2.reset(y_norms2);
+        fvec_norms_L2sqr(y_norms2, y, d, ny);
+        y_norms = y_norms2;
+    }
+
+    for (size_t i0 = 0; i0 < nx; i0 += bs_x) {
+        size_t i1 = i0 + bs_x;
+        if (i1 > nx)
+            i1 = nx;
+
+        res.begin_multiple(i0, i1);
+
+        for (size_t j0 = 0; j0 < ny; j0 += bs_y) {
+            size_t j1 = j0 + bs_y;
+            if (j1 > ny)
+                j1 = ny;
+            /* compute the actual dot products */
+            {
+                float one = 1, zero = 0;
+                FINTEGER nyi = j1 - j0, nxi = i1 - i0, di = d;
+                sgemm_("Transpose",
+                       "Not transpose",
+                       &nyi,
+                       &nxi,
+                       &di,
+                       &one,
+                       y + j0 * d,
+                       &di,
+                       x + i0 * d,
+                       &di,
+                       &zero,
+                       ip_block.get(),
+                       &nyi);
+            }
+#pragma omp parallel for
+            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);
+
+                // constant
+                const __m256 mul_minus2 = _mm256_set1_ps(-2);
+
+                // Track 8 min distances + 8 min indices.
+                // All the distances tracked do not take x_norms[i]
+                //   into account in order to get rid of extra
+                //   _mm256_add_ps(x_norms[i], ...) instructions
+                //   is distance computations.
+                __m256 min_distances =
+                        _mm256_set1_ps(res.dis_tab[i] - x_norms[i]);
+
+                // these indices are local and are relative to j0.
+                // so, value 0 means j0.
+                __m256i min_indices = _mm256_set1_epi32(0);
+
+                __m256i current_indices =
+                        _mm256_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7);
+                const __m256i indices_delta = _mm256_set1_epi32(8);
+
+                // current j index
+                size_t idx_j = 0;
+                size_t count = j1 - j0;
+
+                // 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);
+
+                    // load values for norms
+                    const __m256 y_norm_0 =
+                            _mm256_loadu_ps(y_norms + idx_j + j0 + 0);
+                    const __m256 y_norm_1 =
+                            _mm256_loadu_ps(y_norms + idx_j + j0 + 8);
+
+                    // load values for dot products
+                    const __m256 ip_0 = _mm256_loadu_ps(ip_line + 0);
+                    const __m256 ip_1 = _mm256_loadu_ps(ip_line + 8);
+
+                    // compute dis = y_norm[j] - 2 * dot(x_norm[i], y_norm[j]).
+                    // x_norm[i] was dropped off because it is a constant for a
+                    // given i. We'll deal with it later.
+                    __m256 distances_0 =
+                            _mm256_fmadd_ps(ip_0, mul_minus2, y_norm_0);
+                    __m256 distances_1 =
+                            _mm256_fmadd_ps(ip_1, mul_minus2, y_norm_1);
+
+                    // compare the new distances to the min distances
+                    // for each of the first group of 8 AVX2 components.
+                    const __m256 comparison_0 = _mm256_cmp_ps(
+                            min_distances, distances_0, _CMP_LE_OS);
+
+                    // update min distances and indices with closest vectors if
+                    // needed.
+                    min_distances = _mm256_blendv_ps(
+                            distances_0, min_distances, comparison_0);
+                    min_indices = _mm256_castps_si256(_mm256_blendv_ps(
+                            _mm256_castsi256_ps(current_indices),
+                            _mm256_castsi256_ps(min_indices),
+                            comparison_0));
+                    current_indices =
+                            _mm256_add_epi32(current_indices, indices_delta);
+
+                    // compare the new distances to the min distances
+                    // for each of the second group of 8 AVX2 components.
+                    const __m256 comparison_1 = _mm256_cmp_ps(
+                            min_distances, distances_1, _CMP_LE_OS);
+
+                    // update min distances and indices with closest vectors if
+                    // needed.
+                    min_distances = _mm256_blendv_ps(
+                            distances_1, min_distances, comparison_1);
+                    min_indices = _mm256_castps_si256(_mm256_blendv_ps(
+                            _mm256_castsi256_ps(current_indices),
+                            _mm256_castsi256_ps(min_indices),
+                            comparison_1));
+                    current_indices =
+                            _mm256_add_epi32(current_indices, indices_delta);
+                }
+
+                // dump values and find the minimum distance / minimum index
+                float min_distances_scalar[8];
+                uint32_t min_indices_scalar[8];
+                _mm256_storeu_ps(min_distances_scalar, min_distances);
+                _mm256_storeu_si256(
+                        (__m256i*)(min_indices_scalar), min_indices);
+
+                float current_min_distance = res.dis_tab[i];
+                uint32_t current_min_index = res.ids_tab[i];
+
+                // This unusual comparison is needed to maintain the behavior
+                // of the original implementation: if two indices are
+                // represented with equal distance values, then
+                // the index with the min value is returned.
+                for (size_t jv = 0; jv < 8; jv++) {
+                    // add missing x_norms[i]
+                    float distance_candidate =
+                            min_distances_scalar[jv] + x_norms[i];
+
+                    // negative values can occur for identical vectors
+                    //    due to roundoff errors.
+                    if (distance_candidate < 0)
+                        distance_candidate = 0;
+
+                    int64_t index_candidate = min_indices_scalar[jv] + j0;
+
+                    if (current_min_distance > distance_candidate) {
+                        current_min_distance = distance_candidate;
+                        current_min_index = index_candidate;
+                    } else if (
+                            current_min_distance == distance_candidate &&
+                            current_min_index > index_candidate) {
+                        current_min_index = index_candidate;
+                    }
+                }
+
+                // process leftovers
+                for (; idx_j < count; idx_j++, ip_line++) {
+                    float ip = *ip_line;
+                    float dis = x_norms[i] + y_norms[idx_j + j0] - 2 * ip;
+                    // negative values can occur for identical vectors
+                    //    due to roundoff errors.
+                    if (dis < 0)
+                        dis = 0;
+
+                    if (current_min_distance > dis) {
+                        current_min_distance = dis;
+                        current_min_index = idx_j + j0;
+                    }
+                }
+
+                //
+                res.add_result(i, current_min_distance, current_min_index);
+            }
+        }
+        InterruptCallback::check();
+    }
+}
+#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);
+    }
+}
+
 } // anonymous namespace
 
 /*******************************************************
@@ -313,24 +554,63 @@ void knn_inner_product(
         size_t d,
         size_t nx,
         size_t ny,
-        float_minheap_array_t* ha) {
-    if (ha->k < distance_compute_min_k_reservoir) {
-        HeapResultHandler<CMin<float, int64_t>> res(
-                ha->nh, ha->val, ha->ids, ha->k);
-        if (nx < distance_compute_blas_threshold) {
+        size_t k,
+        float* val,
+        int64_t* ids,
+        const IDSelector* sel) {
+    int64_t imin = 0;
+    if (auto selr = dynamic_cast<const IDSelectorRange*>(sel)) {
+        imin = std::max(selr->imin, int64_t(0));
+        int64_t imax = std::min(selr->imax, int64_t(ny));
+        ny = imax - imin;
+        y += d * imin;
+        sel = nullptr;
+    }
+    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);
+        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 {
-        ReservoirResultHandler<CMin<float, int64_t>> res(
-                ha->nh, ha->val, ha->ids, ha->k);
-        if (nx < distance_compute_blas_threshold) {
-            exhaustive_inner_product_seq(x, y, d, nx, ny, res);
+        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);
         }
     }
+    if (imin != 0) {
+        for (size_t i = 0; i < nx * k; i++) {
+            if (ids[i] >= 0) {
+                ids[i] += imin;
+            }
+        }
+    }
+}
+
+void knn_inner_product(
+        const float* x,
+        const float* y,
+        size_t d,
+        size_t nx,
+        size_t ny,
+        float_minheap_array_t* res,
+        const IDSelector* sel) {
+    FAISS_THROW_IF_NOT(nx == res->nh);
+    knn_inner_product(x, y, d, nx, ny, res->k, res->val, res->ids, sel);
 }
 
 void knn_L2sqr(
@@ -339,28 +619,55 @@ void knn_L2sqr(
         size_t d,
         size_t nx,
         size_t ny,
-        float_maxheap_array_t* ha,
-        const float* y_norm2) {
-    if (ha->k < distance_compute_min_k_reservoir) {
-        HeapResultHandler<CMax<float, int64_t>> res(
-                ha->nh, ha->val, ha->ids, ha->k);
-
-        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);
-        }
+        size_t k,
+        float* vals,
+        int64_t* ids,
+        const float* y_norm2,
+        const IDSelector* sel) {
+    int64_t imin = 0;
+    if (auto selr = dynamic_cast<const IDSelectorRange*>(sel)) {
+        imin = std::max(selr->imin, int64_t(0));
+        int64_t imax = std::min(selr->imax, int64_t(ny));
+        ny = imax - imin;
+        y += d * imin;
+        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);
+        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(
-                ha->nh, ha->val, ha->ids, ha->k);
-        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);
+        ReservoirResultHandler<CMax<float, int64_t>> res(nx, vals, ids, k);
+        knn_L2sqr_select(x, y, d, nx, ny, res, y_norm2, sel);
+    }
+    if (imin != 0) {
+        for (size_t i = 0; i < nx * k; i++) {
+            if (ids[i] >= 0) {
+                ids[i] += imin;
+            }
         }
     }
 }
 
+void knn_L2sqr(
+        const float* x,
+        const float* y,
+        size_t d,
+        size_t nx,
+        size_t ny,
+        float_maxheap_array_t* res,
+        const float* y_norm2,
+        const IDSelector* sel) {
+    FAISS_THROW_IF_NOT(res->nh == nx);
+    knn_L2sqr(x, y, d, nx, ny, res->k, res->val, res->ids, y_norm2, sel);
+}
+
 /***************************************************************************
  * Range search
  ***************************************************************************/
@@ -372,10 +679,14 @@ void range_search_L2sqr(
         size_t nx,
         size_t ny,
         float radius,
-        RangeSearchResult* res) {
-    RangeSearchResultHandler<CMax<float, int64_t>> resh(res, radius);
-    if (nx < distance_compute_blas_threshold) {
-        exhaustive_L2sqr_seq(x, y, d, nx, ny, resh);
+        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);
     }
@@ -388,9 +699,13 @@ void range_search_inner_product(
         size_t nx,
         size_t ny,
         float radius,
-        RangeSearchResult* res) {
-    RangeSearchResultHandler<CMin<float, int64_t>> resh(res, radius);
-    if (nx < distance_compute_blas_threshold) {
+        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);
@@ -488,16 +803,21 @@ void knn_inner_products_by_idx(
         size_t d,
         size_t nx,
         size_t ny,
-        float_minheap_array_t* res) {
-    size_t k = res->k;
+        size_t k,
+        float* res_vals,
+        int64_t* res_ids,
+        int64_t ld_ids) {
+    if (ld_ids < 0) {
+        ld_ids = ny;
+    }
 
-#pragma omp parallel for
+#pragma omp parallel for if (nx > 100)
     for (int64_t i = 0; i < nx; i++) {
         const float* x_ = x + i * d;
-        const int64_t* idsi = ids + i * ny;
+        const int64_t* idsi = ids + i * ld_ids;
         size_t j;
-        float* __restrict simi = res->get_val(i);
-        int64_t* __restrict idxi = res->get_ids(i);
+        float* __restrict simi = res_vals + i * k;
+        int64_t* __restrict idxi = res_ids + i * k;
         minheap_heapify(k, simi, idxi);
 
         for (j = 0; j < ny; j++) {
@@ -520,16 +840,20 @@ void knn_L2sqr_by_idx(
         size_t d,
         size_t nx,
         size_t ny,
-        float_maxheap_array_t* res) {
-    size_t k = res->k;
-
-#pragma omp parallel for
+        size_t k,
+        float* res_vals,
+        int64_t* res_ids,
+        int64_t ld_ids) {
+    if (ld_ids < 0) {
+        ld_ids = ny;
+    }
+#pragma omp parallel for if (nx > 100)
     for (int64_t i = 0; i < nx; i++) {
         const float* x_ = x + i * d;
-        const int64_t* __restrict idsi = ids + i * ny;
-        float* __restrict simi = res->get_val(i);
-        int64_t* __restrict idxi = res->get_ids(i);
-        maxheap_heapify(res->k, simi, idxi);
+        const int64_t* __restrict idsi = ids + i * ld_ids;
+        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++) {
             float disij = fvec_L2sqr(x_, y + d * idsi[j], d);
 
@@ -537,7 +861,7 @@ void knn_L2sqr_by_idx(
                 maxheap_replace_top(k, simi, idxi, disij, idsi[j]);
             }
         }
-        maxheap_reorder(res->k, simi, idxi);
+        maxheap_reorder(k, simi, idxi);
     }
 }