faiss 0.2.6 → 0.2.7

data/vendor/faiss/faiss/utils/approx_topk/approx_topk.h

@@ -0,0 +1,84 @@
+/**
+ * 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.
+ */
+
+// This file contains an implementation of approximate top-k search
+// using heap. It was initially created for a beam search.
+//
+// The core idea is the following.
+// Say we need to find beam_size indices with the minimal distance
+// values. It is done via heap (priority_queue) using the following
+// pseudocode:
+//
+//   def baseline():
+//     distances = np.empty([beam_size * n], dtype=float)
+//     indices = np.empty([beam_size * n], dtype=int)
+//
+//     heap = Heap(max_heap_size=beam_size)
+//
+//     for i in range(0, beam_size * n):
+//         heap.push(distances[i], indices[i])
+//
+// Basically, this is what heap_addn() function from utils/Heap.h does.
+//
+// The following scheme can be used for approximate beam search.
+// Say, we need to find elements with min distance.
+// Basically, we split n elements of every beam into NBUCKETS buckets
+// and track the index with the minimal distance for every bucket.
+// This can be effectively SIMD-ed and significantly lowers the number
+// of operations, but yields approximate results for beam_size >= 2.
+//
+//  def approximate_v1():
+//    distances = np.empty([beam_size * n], dtype=float)
+//    indices = np.empty([beam_size * n], dtype=int)
+//
+//    heap = Heap(max_heap_size=beam_size)
+//
+//    for beam in range(0, beam_size):
+//      # The value of 32 is just an example.
+//      # The value may be varied: the larger the value is,
+//      #  the slower and the more precise vs baseline beam search is
+//      NBUCKETS = 32
+//
+//     local_min_distances = [HUGE_VALF] * NBUCKETS
+//     local_min_indices = [0] * NBUCKETS
+//
+//      for i in range(0, n / NBUCKETS):
+//        for j in range(0, NBUCKETS):
+//          idx = beam * n + i * NBUCKETS + j
+//          if distances[idx] < local_min_distances[j]:
+//            local_min_distances[i] = distances[idx]
+//            local_min_indices[i] = indices[idx]
+//
+//    for j in range(0, NBUCKETS):
+//      heap.push(local_min_distances[j], local_min_indices[j])
+//
+// The accuracy can be improved by tracking min-2 elements for every
+// bucket. Such a min-2 implementation with NBUCKETS buckets provides
+// better accuracy than top-1 implementation with 2 * NBUCKETS buckets.
+// Min-3 is also doable. One can use min-N approach, but I'm not sure
+// whether min-4 and above are practical, because of the lack of SIMD
+// registers (unless AVX-512 version is used).
+//
+// C++ template for top-N implementation is provided. The code
+// assumes that indices[idx] == idx. One can write a code that lifts
+// such an assumption easily.
+//
+// Currently, the code that tracks elements with min distances is implemented
+//    (Max Heap). Min Heap option can be added easily.
+
+#pragma once
+
+#include <faiss/impl/platform_macros.h>
+
+// the list of available modes is in the following file
+#include <faiss/utils/approx_topk/mode.h>
+
+#ifdef __AVX2__
+#include <faiss/utils/approx_topk/avx2-inl.h>
+#else
+#include <faiss/utils/approx_topk/generic.h>
+#endif

data/vendor/faiss/faiss/utils/approx_topk/avx2-inl.h

@@ -0,0 +1,196 @@
+/**
+ * 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 <immintrin.h>
+
+#include <limits>
+
+#include <faiss/impl/FaissAssert.h>
+#include <faiss/utils/Heap.h>
+
+namespace faiss {
+
+template <typename C, uint32_t NBUCKETS, uint32_t N>
+struct HeapWithBuckets {
+    // this case was not implemented yet.
+};
+
+template <uint32_t NBUCKETS, uint32_t N>
+struct HeapWithBuckets<CMax<float, int>, NBUCKETS, N> {
+    static constexpr uint32_t NBUCKETS_8 = NBUCKETS / 8;
+    static_assert(
+            (NBUCKETS) > 0 && ((NBUCKETS % 8) == 0),
+            "Number of buckets needs to be 8, 16, 24, ...");
+
+    static void addn(
+            // number of elements
+            const uint32_t n,
+            // distances. It is assumed to have n elements.
+            const float* const __restrict distances,
+            // number of best elements to keep
+            const uint32_t k,
+            // output distances
+            float* const __restrict bh_val,
+            // output indices, each being within [0, n) range
+            int32_t* const __restrict bh_ids) {
+        // forward a call to bs_addn with 1 beam
+        bs_addn(1, n, distances, k, bh_val, bh_ids);
+    }
+
+    static void bs_addn(
+            // beam_size parameter of Beam Search algorithm
+            const uint32_t beam_size,
+            // number of elements per beam
+            const uint32_t n_per_beam,
+            // distances. It is assumed to have (n_per_beam * beam_size)
+            // elements.
+            const float* const __restrict distances,
+            // number of best elements to keep
+            const uint32_t k,
+            // output distances
+            float* const __restrict bh_val,
+            // output indices, each being within [0, n_per_beam * beam_size)
+            // range
+            int32_t* const __restrict bh_ids) {
+        // // Basically, the function runs beam_size iterations.
+        // // Every iteration NBUCKETS * N elements are added to a regular heap.
+        // // So, maximum number of added elements is beam_size * NBUCKETS * N.
+        // // This number is expected to be less or equal than k.
+        // FAISS_THROW_IF_NOT_FMT(
+        //         beam_size * NBUCKETS * N >= k,
+        //         "Cannot pick %d elements, only %d. "
+        //         "Check the function and template arguments values.",
+        //         k,
+        //         beam_size * NBUCKETS * N);
+
+        using C = CMax<float, int>;
+
+        // main loop
+        for (uint32_t beam_index = 0; beam_index < beam_size; beam_index++) {
+            __m256 min_distances_i[NBUCKETS_8][N];
+            __m256i min_indices_i[NBUCKETS_8][N];
+
+            for (uint32_t j = 0; j < NBUCKETS_8; j++) {
+                for (uint32_t p = 0; p < N; p++) {
+                    min_distances_i[j][p] =
+                            _mm256_set1_ps(std::numeric_limits<float>::max());
+                    min_indices_i[j][p] =
+                            _mm256_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7);
+                }
+            }
+
+            __m256i current_indices = _mm256_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7);
+            __m256i indices_delta = _mm256_set1_epi32(NBUCKETS);
+
+            const uint32_t nb = (n_per_beam / NBUCKETS) * NBUCKETS;
+
+            // put the data into buckets
+            for (uint32_t ip = 0; ip < nb; ip += NBUCKETS) {
+                for (uint32_t j = 0; j < NBUCKETS_8; j++) {
+                    const __m256 distances_reg = _mm256_loadu_ps(
+                            distances + j * 8 + ip + n_per_beam * beam_index);
+
+                    // loop. Compiler should get rid of unneeded ops
+                    __m256 distance_candidate = distances_reg;
+                    __m256i indices_candidate = current_indices;
+
+                    for (uint32_t p = 0; p < N; p++) {
+                        const __m256 comparison = _mm256_cmp_ps(
+                                min_distances_i[j][p],
+                                distance_candidate,
+                                _CMP_LE_OS);
+
+                        // // blend seems to be slower that min
+                        // const __m256 min_distances_new = _mm256_blendv_ps(
+                        //         distance_candidate,
+                        //         min_distances_i[j][p],
+                        //         comparison);
+                        const __m256 min_distances_new = _mm256_min_ps(
+                                distance_candidate, min_distances_i[j][p]);
+                        const __m256i min_indices_new =
+                                _mm256_castps_si256(_mm256_blendv_ps(
+                                        _mm256_castsi256_ps(indices_candidate),
+                                        _mm256_castsi256_ps(
+                                                min_indices_i[j][p]),
+                                        comparison));
+
+                        // // blend seems to be slower that min
+                        // const __m256 max_distances_new = _mm256_blendv_ps(
+                        //         min_distances_i[j][p],
+                        //         distance_candidate,
+                        //         comparison);
+                        const __m256 max_distances_new = _mm256_max_ps(
+                                min_distances_i[j][p], distances_reg);
+                        const __m256i max_indices_new =
+                                _mm256_castps_si256(_mm256_blendv_ps(
+                                        _mm256_castsi256_ps(
+                                                min_indices_i[j][p]),
+                                        _mm256_castsi256_ps(indices_candidate),
+                                        comparison));
+
+                        distance_candidate = max_distances_new;
+                        indices_candidate = max_indices_new;
+
+                        min_distances_i[j][p] = min_distances_new;
+                        min_indices_i[j][p] = min_indices_new;
+                    }
+                }
+
+                current_indices =
+                        _mm256_add_epi32(current_indices, indices_delta);
+            }
+
+            // fix the indices
+            for (uint32_t j = 0; j < NBUCKETS_8; j++) {
+                const __m256i offset =
+                        _mm256_set1_epi32(n_per_beam * beam_index + j * 8);
+                for (uint32_t p = 0; p < N; p++) {
+                    min_indices_i[j][p] =
+                            _mm256_add_epi32(min_indices_i[j][p], offset);
+                }
+            }
+
+            // merge every bucket into the regular heap
+            for (uint32_t p = 0; p < N; p++) {
+                for (uint32_t j = 0; j < NBUCKETS_8; j++) {
+                    int32_t min_indices_scalar[8];
+                    float min_distances_scalar[8];
+
+                    _mm256_storeu_si256(
+                            (__m256i*)min_indices_scalar, min_indices_i[j][p]);
+                    _mm256_storeu_ps(
+                            min_distances_scalar, min_distances_i[j][p]);
+
+                    // this exact way is needed to maintain the order as if the
+                    // input elements were pushed to the heap sequentially
+                    for (size_t j8 = 0; j8 < 8; j8++) {
+                        const auto value = min_distances_scalar[j8];
+                        const auto index = min_indices_scalar[j8];
+                        if (C::cmp2(bh_val[0], value, bh_ids[0], index)) {
+                            heap_replace_top<C>(
+                                    k, bh_val, bh_ids, value, index);
+                        }
+                    }
+                }
+            }
+
+            // process leftovers
+            for (uint32_t ip = nb; ip < n_per_beam; ip++) {
+                const int32_t index = ip + n_per_beam * beam_index;
+                const float value = distances[index];
+
+                if (C::cmp(bh_val[0], value)) {
+                    heap_replace_top<C>(k, bh_val, bh_ids, value, index);
+                }
+            }
+        }
+    }
+};
+
+} // namespace faiss

data/vendor/faiss/faiss/utils/approx_topk/generic.h

@@ -0,0 +1,138 @@
+/**
+ * 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 <algorithm>
+#include <limits>
+#include <utility>
+
+#include <faiss/impl/FaissAssert.h>
+#include <faiss/utils/Heap.h>
+
+namespace faiss {
+
+// This is the implementation of the idea and it is very slow,
+// because a compiler is unable to vectorize it properly.
+
+template <typename C, uint32_t NBUCKETS, uint32_t N>
+struct HeapWithBuckets {
+    // this case was not implemented yet.
+};
+
+template <uint32_t NBUCKETS, uint32_t N>
+struct HeapWithBuckets<CMax<float, int>, NBUCKETS, N> {
+    static void addn(
+            // number of elements
+            const uint32_t n,
+            // distances. It is assumed to have n elements.
+            const float* const __restrict distances,
+            // number of best elements to keep
+            const uint32_t k,
+            // output distances
+            float* const __restrict bh_val,
+            // output indices, each being within [0, n) range
+            int32_t* const __restrict bh_ids) {
+        // forward a call to bs_addn with 1 beam
+        bs_addn(1, n, distances, k, bh_val, bh_ids);
+    }
+
+    static void bs_addn(
+            // beam_size parameter of Beam Search algorithm
+            const uint32_t beam_size,
+            // number of elements per beam
+            const uint32_t n_per_beam,
+            // distances. It is assumed to have (n_per_beam * beam_size)
+            // elements.
+            const float* const __restrict distances,
+            // number of best elements to keep
+            const uint32_t k,
+            // output distances
+            float* const __restrict bh_val,
+            // output indices, each being within [0, n_per_beam * beam_size)
+            // range
+            int32_t* const __restrict bh_ids) {
+        // // Basically, the function runs beam_size iterations.
+        // // Every iteration NBUCKETS * N elements are added to a regular heap.
+        // // So, maximum number of added elements is beam_size * NBUCKETS * N.
+        // // This number is expected to be less or equal than k.
+        // FAISS_THROW_IF_NOT_FMT(
+        //         beam_size * NBUCKETS * N >= k,
+        //         "Cannot pick %d elements, only %d. "
+        //         "Check the function and template arguments values.",
+        //         k,
+        //         beam_size * NBUCKETS * N);
+
+        using C = CMax<float, int>;
+
+        // main loop
+        for (uint32_t beam_index = 0; beam_index < beam_size; beam_index++) {
+            float min_distances_i[N][NBUCKETS];
+            int min_indices_i[N][NBUCKETS];
+
+            for (uint32_t p = 0; p < N; p++) {
+                for (uint32_t j = 0; j < NBUCKETS; j++) {
+                    min_distances_i[p][j] = std::numeric_limits<float>::max();
+                    min_indices_i[p][j] = 0;
+                }
+            }
+
+            const uint32_t nb = (n_per_beam / NBUCKETS) * NBUCKETS;
+
+            // put the data into buckets
+            for (uint32_t ip = 0; ip < nb; ip += NBUCKETS) {
+                for (uint32_t j = 0; j < NBUCKETS; j++) {
+                    const int index = j + ip + n_per_beam * beam_index;
+                    const float distance = distances[index];
+
+                    int index_candidate = index;
+                    float distance_candidate = distance;
+
+                    for (uint32_t p = 0; p < N; p++) {
+                        if (distance_candidate < min_distances_i[p][j]) {
+                            std::swap(
+                                    distance_candidate, min_distances_i[p][j]);
+                            std::swap(index_candidate, min_indices_i[p][j]);
+                        }
+                    }
+                }
+            }
+
+            // merge every bucket into the regular heap
+            for (uint32_t p = 0; p < N; p++) {
+                for (uint32_t j = 0; j < NBUCKETS; j++) {
+                    // this exact way is needed to maintain the order as if the
+                    // input elements were pushed to the heap sequentially
+
+                    if (C::cmp2(bh_val[0],
+                                min_distances_i[p][j],
+                                bh_ids[0],
+                                min_indices_i[p][j])) {
+                        heap_replace_top<C>(
+                                k,
+                                bh_val,
+                                bh_ids,
+                                min_distances_i[p][j],
+                                min_indices_i[p][j]);
+                    }
+                }
+            }
+
+            // process leftovers
+            for (uint32_t ip = nb; ip < n_per_beam; ip++) {
+                const int32_t index = ip + n_per_beam * beam_index;
+                const float value = distances[index];
+
+                if (C::cmp(bh_val[0], value)) {
+                    heap_replace_top<C>(k, bh_val, bh_ids, value, index);
+                }
+            }
+        }
+    }
+};
+
+} // namespace faiss

data/vendor/faiss/faiss/utils/approx_topk/mode.h

@@ -0,0 +1,34 @@
+/**
+ * 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
+
+/// Represents the mode of use of approximate top-k computations
+/// that allows to trade accuracy vs speed. So, every options
+/// besides EXACT_TOPK increases the speed.
+///
+/// B represents the number of buckets.
+/// D is the number of min-k elements to track within every bucket.
+///
+/// Default option is EXACT_TOPK.
+/// APPROX_TOPK_BUCKETS_B16_D2 is worth starting from, if you'd like
+/// to experiment a bit.
+///
+/// It seems that only the limited number of combinations are
+/// meaningful, because of the limited supply of SIMD registers.
+/// Also, certain combinations, such as B32_D1 and B16_D1, were concluded
+/// to be not very precise in benchmarks, so ones were not introduced.
+///
+/// TODO: Consider d-ary SIMD heap.
+
+enum ApproxTopK_mode_t : int {
+    EXACT_TOPK = 0,
+    APPROX_TOPK_BUCKETS_B32_D2 = 1,
+    APPROX_TOPK_BUCKETS_B8_D3 = 2,
+    APPROX_TOPK_BUCKETS_B16_D2 = 3,
+    APPROX_TOPK_BUCKETS_B8_D2 = 4,
+};