faiss 0.5.3 → 0.6.1

data/vendor/faiss/faiss/impl/ScalarQuantizer.h

@@ -5,8 +5,6 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-// -*- c++ -*-
-
 #pragma once
 
 #include <faiss/impl/AuxIndexStructures.h>
@@ -35,6 +33,13 @@ struct ScalarQuantizer : Quantizer {
         QT_bf16,
         QT_8bit_direct_signed, ///< fast indexing of signed int8s ranging from
                                ///< [-128 to 127]
+        QT_0bit, ///< 0 bits per component, centroid-only distance (for IVF)
+        QT_1bit_tqmse, ///< TurboQuant MSE-optimized, 1 bit per component
+        QT_2bit_tqmse, ///< TurboQuant MSE-optimized, 2 bits per component
+        QT_3bit_tqmse, ///< TurboQuant MSE-optimized, 3 bits per component
+        QT_4bit_tqmse, ///< TurboQuant MSE-optimized, 4 bits per component
+        QT_8bit_tqmse, ///< TurboQuant MSE-optimized, 8 bits per component
+        QT_count
     };
 
     QuantizerType qtype = QT_8bit;
@@ -60,7 +65,7 @@ struct ScalarQuantizer : Quantizer {
     /// trained values (including the range)
     std::vector<float> trained;
 
-    ScalarQuantizer(size_t d, QuantizerType qtype);
+    ScalarQuantizer(size_t d_in, QuantizerType qtype_in);
     ScalarQuantizer();
 
     /// updates internal values based on qtype and d
@@ -102,6 +107,25 @@ struct ScalarQuantizer : Quantizer {
 
         virtual float query_to_code(const uint8_t* code) const = 0;
 
+        /// Compute four query-to-code distances in one call. Default loops
+        /// query_to_code four times; per-SIMD specializations may batch the
+        /// inner dim loop across the four codes to amortize query state and
+        /// expose ILP across independent accumulators.
+        virtual void query_to_codes_batch_4(
+                const uint8_t* code_0,
+                const uint8_t* code_1,
+                const uint8_t* code_2,
+                const uint8_t* code_3,
+                float& dis0,
+                float& dis1,
+                float& dis2,
+                float& dis3) const {
+            dis0 = query_to_code(code_0);
+            dis1 = query_to_code(code_1);
+            dis2 = query_to_code(code_2);
+            dis3 = query_to_code(code_3);
+        }
+
         float distance_to_code(const uint8_t* code) final {
             return query_to_code(code);
         }

data/vendor/faiss/faiss/impl/ThreadedIndex-inl.h

@@ -19,8 +19,8 @@ ThreadedIndex<IndexT>::ThreadedIndex(bool threaded)
         : ThreadedIndex(0, threaded) {}
 
 template <typename IndexT>
-ThreadedIndex<IndexT>::ThreadedIndex(int d, bool threaded)
-        : IndexT(d), isThreaded_(threaded) {}
+ThreadedIndex<IndexT>::ThreadedIndex(int d_in, bool threaded)
+        : IndexT(d_in), isThreaded_(threaded) {}
 
 template <typename IndexT>
 ThreadedIndex<IndexT>::~ThreadedIndex() {
@@ -122,11 +122,12 @@ void ThreadedIndex<IndexT>::runOnIndex(std::function<void(int, IndexT*)> f) {
     if (isThreaded_) {
         std::vector<std::future<bool>> v;
 
-        for (int i = 0; i < this->indices_.size(); ++i) {
+        for (size_t i = 0; i < this->indices_.size(); ++i) {
             auto& p = this->indices_[i];
             auto indexPtr = p.first;
+            int idx = static_cast<int>(i);
             v.emplace_back(
-                    p.second->add([f, i, indexPtr]() { f(i, indexPtr); }));
+                    p.second->add([f, idx, indexPtr]() { f(idx, indexPtr); }));
         }
 
         waitAndHandleFutures(v);
@@ -135,13 +136,14 @@ void ThreadedIndex<IndexT>::runOnIndex(std::function<void(int, IndexT*)> f) {
         // while letting everything else run to completion
         std::vector<std::pair<int, std::exception_ptr>> exceptions;
 
-        for (int i = 0; i < this->indices_.size(); ++i) {
+        for (size_t i = 0; i < this->indices_.size(); ++i) {
             auto& p = this->indices_[i];
             try {
-                f(i, p.first);
+                f(static_cast<int>(i), p.first);
             } catch (...) {
                 exceptions.emplace_back(
-                        std::make_pair(i, std::current_exception()));
+                        std::make_pair(
+                                static_cast<int>(i), std::current_exception()));
             }
         }
 
@@ -164,10 +166,10 @@ void ThreadedIndex<IndexT>::reset() {
 }
 
 template <typename IndexT>
-void ThreadedIndex<IndexT>::onAfterAddIndex(IndexT* index) {}
+void ThreadedIndex<IndexT>::onAfterAddIndex(IndexT* /* index */) {}
 
 template <typename IndexT>
-void ThreadedIndex<IndexT>::onAfterRemoveIndex(IndexT* index) {}
+void ThreadedIndex<IndexT>::onAfterRemoveIndex(IndexT* /* index */) {}
 
 template <typename IndexT>
 void ThreadedIndex<IndexT>::waitAndHandleFutures(
@@ -176,14 +178,15 @@ void ThreadedIndex<IndexT>::waitAndHandleFutures(
     // exceptions that are generated
     std::vector<std::pair<int, std::exception_ptr>> exceptions;
 
-    for (int i = 0; i < v.size(); ++i) {
+    for (size_t i = 0; i < v.size(); ++i) {
         auto& fut = v[i];
 
         try {
             fut.get();
         } catch (...) {
             exceptions.emplace_back(
-                    std::make_pair(i, std::current_exception()));
+                    std::make_pair(
+                            static_cast<int>(i), std::current_exception()));
         }
     }
 

data/vendor/faiss/faiss/impl/VisitedTable.cpp

@@ -0,0 +1,42 @@
+/*
+ * 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.
+ */
+
+#include <faiss/impl/VisitedTable.h>
+
+#include <cstring>
+
+namespace faiss {
+
+// The vector strategy is faster for get()/set(), but O(size) to initialize.
+// advance() is O(1) except every 250 calls, which are O(size).
+// The hash set strategy is a constant factor slower for get()/set(),
+// but O(1) to construct and O(visits) to advance.
+// A size of ~1M seems to be the threshold where the hash set wins.
+size_t visited_table_hashset_threshold = 500000;
+
+VisitedTable::VisitedTable(size_t size, std::optional<bool> use_hashset)
+        : visno(use_hashset.value_or(size >= visited_table_hashset_threshold)
+                        ? 0
+                        : 1) {
+    if (visno != 0) {
+        visited.resize(size, 0);
+    }
+}
+
+void VisitedTable::advance() {
+    if (visno == 0) {
+        visited_set.clear();
+    } else if (visno < 254) {
+        // 254 rather than 255 because sometimes we use visno and visno+1
+        ++visno;
+    } else {
+        memset(visited.data(), 0, sizeof(visited[0]) * visited.size());
+        visno = 1;
+    }
+}
+
+} // namespace faiss

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

@@ -0,0 +1,76 @@
+/*
+ * 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.
+ */
+
+#ifndef FAISS_VISITED_TABLE_H
+#define FAISS_VISITED_TABLE_H
+
+#include <stdint.h>
+
+#include <optional>
+#include <unordered_set>
+#include <vector>
+
+#include <faiss/impl/platform_macros.h>
+#include <faiss/utils/prefetch.h>
+
+namespace faiss {
+
+FAISS_API extern size_t visited_table_hashset_threshold;
+
+/// 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.
+
+    // If use_hashset is nullopt, the use of a hashset will be determined by
+    // size >= visited_table_hashset_threshold.
+    explicit VisitedTable(
+            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) {
+            return false;
+        } else {
+            visited[no] = visno;
+            return true;
+        }
+    }
+
+    /// pre-allocate bucket space to avoid rehashing during repeated set() calls
+    void reserve(size_t n) {
+        if (visno == 0) {
+            visited_set.reserve(n);
+        }
+    }
+
+    /// get flag #no
+    bool get(size_t no) const {
+        if (visno == 0) {
+            return visited_set.count(no) != 0;
+        } else {
+            return visited[no] == visno;
+        }
+    }
+
+    void prefetch(size_t no) const {
+        if (visno != 0) {
+            prefetch_L2(&visited[no]);
+        }
+    }
+
+    /// reset all flags to false
+    void advance();
+};
+
+} // namespace faiss
+
+#endif

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