faiss 0.6.0 → 0.6.2

data/vendor/faiss/faiss/impl/VisitedTable.h

@@ -10,6 +10,7 @@
 
 #include <stdint.h>
 
+#include <memory>
 #include <optional>
 #include <unordered_set>
 #include <vector>
@@ -21,47 +22,88 @@ namespace faiss {
 
 FAISS_API extern size_t visited_table_hashset_threshold;
 
-/// A fast, reusable Visited Set for graph search algorithms.
+/// Abstract base class for a fast, reusable Visited Set for graph search
+/// algorithms.
 struct VisitedTable {
-    std::vector<uint8_t> visited;
-    std::unordered_set<size_t> visited_set;
-    uint8_t visno; // 0 if using visited_set, 1..250 if using vector.
+    virtual ~VisitedTable() = default;
 
-    // If use_hashset is nullopt, the use of a hashset will be determined by
-    // size >= visited_table_hashset_threshold.
-    explicit VisitedTable(
+    /// set flag #no to true, return whether this changed it.
+    virtual bool set(size_t no) = 0;
+
+    /// get flag #no
+    virtual bool get(size_t no) const = 0;
+
+    /// prefetch flag #no
+    virtual void prefetch(size_t no) const = 0;
+
+    /// pre-allocate bucket space to avoid rehashing during repeated set() calls
+    virtual void reserve(size_t /*n*/) {}
+
+    /// reset all flags to false
+    virtual void advance() = 0;
+
+    /// Factory method to create appropriate implementation.
+    /// If use_hashset is nullopt, the use of a hashset will be determined by
+    /// size >= visited_table_hashset_threshold.
+    static std::unique_ptr<VisitedTable> create(
             size_t size,
             std::optional<bool> use_hashset = std::nullopt);
+};
 
-    /// set flag #no to true, return whether this changed it.
-    bool set(size_t no) {
-        if (visno == 0) {
-            return visited_set.insert(no).second;
-        } else if (visited[no] == visno) {
+/// Set-based implementation using unordered_set.
+/// O(1) to construct and O(visits) to advance.
+struct VisitedTableSet FAISS_FINAL : VisitedTable {
+    std::unordered_set<size_t> visited_set;
+
+    VisitedTableSet() = default;
+
+    bool set(size_t no) final {
+        return visited_set.insert(no).second;
+    }
+
+    bool get(size_t no) const final {
+        return visited_set.count(no) != 0;
+    }
+
+    void prefetch(size_t /*no*/) const final {
+        // No-op for set-based implementation
+    }
+
+    void reserve(size_t n) final {
+        visited_set.reserve(n);
+    }
+
+    void advance() final {
+        visited_set.clear();
+    }
+};
+
+/// Vector-based implementation using a versioned byte array.
+/// Faster for get()/set(), but O(size) to initialize.
+/// advance() is O(1) except every 250 calls, which are O(size).
+struct VisitedTableVector FAISS_FINAL : VisitedTable {
+    std::vector<uint8_t> visited;
+    uint8_t visno{1}; // Version number, 1..254
+
+    explicit VisitedTableVector(size_t size) : visited(size, 0) {}
+
+    bool set(size_t no) final {
+        if (visited[no] == visno) {
             return false;
-        } else {
-            visited[no] = visno;
-            return true;
         }
+        visited[no] = visno;
+        return true;
     }
 
-    /// get flag #no
-    bool get(size_t no) const {
-        if (visno == 0) {
-            return visited_set.count(no) != 0;
-        } else {
-            return visited[no] == visno;
-        }
+    bool get(size_t no) const final {
+        return visited[no] == visno;
     }
 
-    void prefetch(size_t no) const {
-        if (visno != 0) {
-            prefetch_L2(&visited[no]);
-        }
+    void prefetch(size_t no) const final {
+        prefetch_L2(&visited[no]);
     }
 
-    /// reset all flags to false
-    void advance();
+    void advance() final;
 };
 
 } // namespace faiss

data/vendor/faiss/faiss/impl/approx_topk/approx_topk.h

@@ -0,0 +1,276 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// Approximate top-k search using bucketed heaps with SIMD acceleration.
+//
+// HeapWithBuckets<C, NBUCKETS, N> splits n elements into NBUCKETS buckets,
+// tracks the top-N per bucket using SIMD, then merges into a regular heap.
+// This trades a small accuracy loss for significant speedup.
+//
+// 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[j] = distances[idx]
+//            local_min_indices[j] = 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.
+//
+// Dispatch:
+//   AVX2 / ARM_NEON  →  HeapWithBucketsCMaxFloat (simdlib256-inl.h)
+//   NONE (scalar)    →  HeapWithBucketsGenericCMaxFloat (generic.h)
+//
+// The SIMD definitions live in simdlib256-inl.h (only included by per-ISA
+// .cpp files). Common TUs see only declarations here, so no extern template
+// suppression is needed.
+
+#pragma once
+
+#include <cstdint>
+
+#include <faiss/impl/approx_topk/generic.h>
+#include <faiss/impl/platform_macros.h>
+#include <faiss/impl/simd_dispatch.h>
+#include <faiss/utils/Heap.h>
+#include <faiss/utils/simd_levels.h>
+
+// -----------------------------------------------------------------------
+// ApproxTopK_mode_t
+// -----------------------------------------------------------------------
+
+/// 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 they were not introduced.
+
+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,
+};
+
+// -----------------------------------------------------------------------
+// HeapWithBuckets dispatch
+// -----------------------------------------------------------------------
+
+namespace faiss {
+
+// Primary template: declared, bs_addn NOT defined here.
+// The out-of-line definition lives in simdlib256-inl.h, which is only
+// included by the per-ISA .cpp files (avx2.cpp, neon.cpp).
+template <uint32_t NBUCKETS, uint32_t N, SIMDLevel SL>
+struct HeapWithBucketsCMaxFloat {
+    static_assert(
+            (NBUCKETS) > 0 && ((NBUCKETS % 8) == 0),
+            "Number of buckets needs to be 8, 16, 24, ...");
+
+    static void addn(
+            const uint32_t n,
+            const float* const __restrict distances,
+            const uint32_t k,
+            float* const __restrict bh_val,
+            int32_t* const __restrict bh_ids) {
+        bs_addn(1, n, distances, k, bh_val, bh_ids);
+    }
+
+    // Declared but not defined — resolved at link time from avx2.cpp/neon.cpp.
+    static void bs_addn(
+            const uint32_t beam_size,
+            const uint32_t n_per_beam,
+            const float* const __restrict distances,
+            const uint32_t k,
+            float* const __restrict bh_val,
+            int32_t* const __restrict bh_ids);
+};
+
+// NONE specialization: delegates to the scalar generic implementation.
+template <uint32_t NBUCKETS, uint32_t N>
+struct HeapWithBucketsCMaxFloat<NBUCKETS, N, SIMDLevel::NONE> {
+    static void addn(
+            const uint32_t n,
+            const float* const __restrict distances,
+            const uint32_t k,
+            float* const __restrict bh_val,
+            int32_t* const __restrict bh_ids) {
+        bs_addn(1, n, distances, k, bh_val, bh_ids);
+    }
+
+    static void bs_addn(
+            const uint32_t beam_size,
+            const uint32_t n_per_beam,
+            const float* const __restrict distances,
+            const uint32_t k,
+            float* const __restrict bh_val,
+            int32_t* const __restrict bh_ids) {
+        HeapWithBucketsGenericCMaxFloat<NBUCKETS, N>::bs_addn(
+                beam_size, n_per_beam, distances, k, bh_val, bh_ids);
+    }
+};
+
+// Primary template — not implemented for arbitrary comparators.
+template <typename C, uint32_t NBUCKETS, uint32_t N>
+struct HeapWithBuckets {
+    static_assert(
+            sizeof(C) == 0,
+            "HeapWithBuckets: unsupported comparator type");
+};
+
+// Partial specialization for CMax<float, int> that dispatches
+// via with_simd_level_256bit.
+template <uint32_t NBUCKETS, uint32_t N>
+struct HeapWithBuckets<CMax<float, int>, NBUCKETS, N> {
+    static void addn(
+            const uint32_t n,
+            const float* const __restrict distances,
+            const uint32_t k,
+            float* const __restrict bh_val,
+            int32_t* const __restrict bh_ids) {
+        bs_addn(1, n, distances, k, bh_val, bh_ids);
+    }
+
+    static void bs_addn(
+            const uint32_t beam_size,
+            const uint32_t n_per_beam,
+            const float* const __restrict distances,
+            const uint32_t k,
+            float* const __restrict bh_val,
+            int32_t* const __restrict bh_ids) {
+        with_simd_level_256bit([&]<SIMDLevel SL>() {
+            HeapWithBucketsCMaxFloat<NBUCKETS, N, SL>::bs_addn(
+                    beam_size, n_per_beam, distances, k, bh_val, bh_ids);
+        });
+    }
+};
+
+// -----------------------------------------------------------------------
+// approx_topk_by_mode: consolidates the mode switch + dispatch pattern
+// used by residual_quantizer_encode_steps.cpp and other callers.
+// -----------------------------------------------------------------------
+
+// SL-parameterized version for callers that have already resolved the
+// SIMD level (e.g., inside a with_simd_level_256bit lambda).
+template <SIMDLevel SL>
+inline void approx_topk_by_mode(
+        ApproxTopK_mode_t mode,
+        uint32_t beam_size,
+        uint32_t n_per_beam,
+        const float* distances,
+        uint32_t k,
+        float* bh_val,
+        int32_t* bh_ids) {
+    using C = CMax<float, int>;
+    auto approx = [&]<uint32_t NB, uint32_t ND>() {
+        HeapWithBucketsCMaxFloat<NB, ND, SL>::bs_addn(
+                beam_size, n_per_beam, distances, k, bh_val, bh_ids);
+    };
+    switch (mode) {
+        case ApproxTopK_mode_t::APPROX_TOPK_BUCKETS_B8_D3:
+            approx.template operator()<8, 3>();
+            break;
+        case ApproxTopK_mode_t::APPROX_TOPK_BUCKETS_B8_D2:
+            approx.template operator()<8, 2>();
+            break;
+        case ApproxTopK_mode_t::APPROX_TOPK_BUCKETS_B16_D2:
+            approx.template operator()<16, 2>();
+            break;
+        case ApproxTopK_mode_t::APPROX_TOPK_BUCKETS_B32_D2:
+            approx.template operator()<32, 2>();
+            break;
+        default:
+            heap_addn<C>(
+                    k,
+                    bh_val,
+                    bh_ids,
+                    distances,
+                    nullptr,
+                    beam_size * n_per_beam);
+            break;
+    }
+}
+
+// Non-SL wrapper that dispatches via with_simd_level_256bit.
+inline void approx_topk_by_mode(
+        ApproxTopK_mode_t mode,
+        uint32_t beam_size,
+        uint32_t n_per_beam,
+        const float* distances,
+        uint32_t k,
+        float* bh_val,
+        int32_t* bh_ids) {
+    with_simd_level_256bit([&]<SIMDLevel SL>() {
+        approx_topk_by_mode<SL>(
+                mode, beam_size, n_per_beam, distances, k, bh_val, bh_ids);
+    });
+}
+
+} // namespace faiss

data/vendor/faiss/faiss/impl/approx_topk/avx2.cpp

@@ -0,0 +1,68 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// Explicit template instantiations of HeapWithBucketsCMaxFloat and
+// accum_and_*_tab for SIMDLevel::AVX2.
+
+#ifdef COMPILE_SIMD_AVX2
+
+#include <faiss/impl/approx_topk/rq_beam_search_tab-inl.h>
+#include <faiss/impl/approx_topk/simdlib256-inl.h>
+#include <faiss/impl/simdlib/simdlib_avx2.h>
+
+namespace faiss {
+
+template struct HeapWithBucketsCMaxFloat<8, 3, SIMDLevel::AVX2>;
+template struct HeapWithBucketsCMaxFloat<8, 2, SIMDLevel::AVX2>;
+template struct HeapWithBucketsCMaxFloat<16, 2, SIMDLevel::AVX2>;
+template struct HeapWithBucketsCMaxFloat<16, 1, SIMDLevel::AVX2>;
+template struct HeapWithBucketsCMaxFloat<32, 2, SIMDLevel::AVX2>;
+
+#define INSTANTIATE_ACCUM_TAB(M)                                 \
+    template void accum_and_store_tab<M, 4, SIMDLevel::AVX2>(    \
+            size_t,                                              \
+            const float* __restrict,                             \
+            const uint64_t* __restrict,                          \
+            const int32_t* __restrict,                           \
+            size_t,                                              \
+            size_t,                                              \
+            size_t,                                              \
+            float* __restrict);                                  \
+    template void accum_and_add_tab<M, 4, SIMDLevel::AVX2>(      \
+            size_t,                                              \
+            const float* __restrict,                             \
+            const uint64_t* __restrict,                          \
+            const int32_t* __restrict,                           \
+            size_t,                                              \
+            size_t,                                              \
+            size_t,                                              \
+            float* __restrict);                                  \
+    template void accum_and_finalize_tab<M, 4, SIMDLevel::AVX2>( \
+            const float* __restrict,                             \
+            const uint64_t* __restrict,                          \
+            const int32_t* __restrict,                           \
+            size_t,                                              \
+            size_t,                                              \
+            size_t,                                              \
+            const float* __restrict,                             \
+            const float* __restrict,                             \
+            float* __restrict);
+
+INSTANTIATE_ACCUM_TAB(1)
+INSTANTIATE_ACCUM_TAB(2)
+INSTANTIATE_ACCUM_TAB(3)
+INSTANTIATE_ACCUM_TAB(4)
+INSTANTIATE_ACCUM_TAB(5)
+INSTANTIATE_ACCUM_TAB(6)
+INSTANTIATE_ACCUM_TAB(7)
+INSTANTIATE_ACCUM_TAB(8)
+
+#undef INSTANTIATE_ACCUM_TAB
+
+} // namespace faiss
+
+#endif // COMPILE_SIMD_AVX2

data/vendor/faiss/faiss/{utils → impl}/approx_topk/generic.h

@@ -16,16 +16,12 @@
 
 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.
-};
+// Scalar (generic) implementation of HeapWithBuckets.
+// This is correct but slow because a compiler is unable to
+// vectorize it properly. Used as the SIMDLevel::NONE fallback.
 
 template <uint32_t NBUCKETS, uint32_t N>
-struct HeapWithBuckets<CMax<float, int>, NBUCKETS, N> {
+struct HeapWithBucketsGenericCMaxFloat {
     static void addn(
             // number of elements
             const uint32_t n,
@@ -135,4 +131,15 @@ struct HeapWithBuckets<CMax<float, int>, NBUCKETS, N> {
     }
 };
 
+// Legacy name kept for backward compatibility (used when
+// approx_topk.h is not included).
+template <typename C, uint32_t NBUCKETS, uint32_t N>
+struct HeapWithBucketsGeneric {
+    // not implemented for arbitrary C
+};
+
+template <uint32_t NBUCKETS, uint32_t N>
+struct HeapWithBucketsGeneric<CMax<float, int>, NBUCKETS, N>
+        : HeapWithBucketsGenericCMaxFloat<NBUCKETS, N> {};
+
 } // namespace faiss

data/vendor/faiss/faiss/impl/approx_topk/neon.cpp

@@ -0,0 +1,68 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// Explicit template instantiations of HeapWithBucketsCMaxFloat and
+// accum_and_*_tab for SIMDLevel::ARM_NEON.
+
+#ifdef COMPILE_SIMD_ARM_NEON
+
+#include <faiss/impl/approx_topk/rq_beam_search_tab-inl.h>
+#include <faiss/impl/approx_topk/simdlib256-inl.h>
+#include <faiss/impl/simdlib/simdlib_neon.h>
+
+namespace faiss {
+
+template struct HeapWithBucketsCMaxFloat<8, 3, SIMDLevel::ARM_NEON>;
+template struct HeapWithBucketsCMaxFloat<8, 2, SIMDLevel::ARM_NEON>;
+template struct HeapWithBucketsCMaxFloat<16, 2, SIMDLevel::ARM_NEON>;
+template struct HeapWithBucketsCMaxFloat<16, 1, SIMDLevel::ARM_NEON>;
+template struct HeapWithBucketsCMaxFloat<32, 2, SIMDLevel::ARM_NEON>;
+
+#define INSTANTIATE_ACCUM_TAB(M)                                     \
+    template void accum_and_store_tab<M, 4, SIMDLevel::ARM_NEON>(    \
+            size_t,                                                  \
+            const float* __restrict,                                 \
+            const uint64_t* __restrict,                              \
+            const int32_t* __restrict,                               \
+            size_t,                                                  \
+            size_t,                                                  \
+            size_t,                                                  \
+            float* __restrict);                                      \
+    template void accum_and_add_tab<M, 4, SIMDLevel::ARM_NEON>(      \
+            size_t,                                                  \
+            const float* __restrict,                                 \
+            const uint64_t* __restrict,                              \
+            const int32_t* __restrict,                               \
+            size_t,                                                  \
+            size_t,                                                  \
+            size_t,                                                  \
+            float* __restrict);                                      \
+    template void accum_and_finalize_tab<M, 4, SIMDLevel::ARM_NEON>( \
+            const float* __restrict,                                 \
+            const uint64_t* __restrict,                              \
+            const int32_t* __restrict,                               \
+            size_t,                                                  \
+            size_t,                                                  \
+            size_t,                                                  \
+            const float* __restrict,                                 \
+            const float* __restrict,                                 \
+            float* __restrict);
+
+INSTANTIATE_ACCUM_TAB(1)
+INSTANTIATE_ACCUM_TAB(2)
+INSTANTIATE_ACCUM_TAB(3)
+INSTANTIATE_ACCUM_TAB(4)
+INSTANTIATE_ACCUM_TAB(5)
+INSTANTIATE_ACCUM_TAB(6)
+INSTANTIATE_ACCUM_TAB(7)
+INSTANTIATE_ACCUM_TAB(8)
+
+#undef INSTANTIATE_ACCUM_TAB
+
+} // namespace faiss
+
+#endif // COMPILE_SIMD_ARM_NEON