faiss 0.6.0 → 0.6.2

data/vendor/faiss/faiss/impl/hnsw/LockVector.cpp

@@ -0,0 +1,54 @@
+/*
+ * 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/hnsw/LockVector.h>
+
+#include <cstdlib>
+
+#include <utility>
+
+namespace faiss {
+
+LockVector::LockVector(LockVector&& other) noexcept
+        : data_(other.data_), size_(other.size_), capacity_(other.capacity_) {
+    other.data_ = nullptr;
+    other.size_ = 0;
+    other.capacity_ = 0;
+}
+
+void LockVector::prepare(size_t new_size) {
+    if (new_size <= size_) {
+        return;
+    }
+    if (new_size > capacity_) {
+        // Ensure geometric capacity growth.
+        size_t new_cap = std::max(new_size, capacity_ * 2);
+        // Just destroy old and init fresh; omp_lock_t is not copyable.
+        clear();
+        data_ = static_cast<omp_lock_t*>(malloc(new_cap * sizeof(omp_lock_t)));
+        FAISS_THROW_IF_NOT(data_ != nullptr);
+        capacity_ = new_cap;
+    }
+    for (size_t i = size_; i < new_size; i++) {
+        omp_init_lock(&data_[i]);
+    }
+    size_ = new_size;
+}
+
+void LockVector::clear() {
+    if (data_) {
+        for (size_t i = 0; i < size_; i++) {
+            omp_destroy_lock(&data_[i]);
+        }
+        free(data_);
+        data_ = nullptr;
+    }
+    size_ = 0;
+    capacity_ = 0;
+}
+
+} // namespace faiss

data/vendor/faiss/faiss/impl/hnsw/LockVector.h

@@ -0,0 +1,64 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include <omp.h>
+
+#include <faiss/impl/FaissAssert.h>
+
+namespace faiss {
+
+/// Contiguous, growable array of locks with geometric growth.
+class LockVector {
+   public:
+    LockVector() = default;
+    explicit LockVector(size_t n) {
+        prepare(n);
+    }
+    // Copy ctor for clone(), initialized as empty.
+    LockVector(const LockVector&) : LockVector() {}
+    LockVector(LockVector&& other) noexcept;
+
+    ~LockVector() {
+        clear();
+    }
+
+    LockVector& operator=(const LockVector&) = delete;
+    LockVector& operator=(LockVector&& other) = delete;
+
+    size_t size() const {
+        return size_;
+    }
+
+    // Ensure size is at least 'new_size'. No locks may be held.
+    void prepare(size_t new_size);
+    // Release all locks and free memory. No locks may be held.
+    void clear();
+
+    void lock(size_t i) {
+        FAISS_CHECK_RANGE_DEBUG(i, 0, size_);
+        omp_set_lock(&data_[i]);
+    }
+
+    void unlock(size_t i) {
+        FAISS_CHECK_RANGE_DEBUG(i, 0, size_);
+        omp_unset_lock(&data_[i]);
+    }
+
+    bool try_lock(size_t i) {
+        FAISS_CHECK_RANGE_DEBUG(i, 0, size_);
+        return omp_test_lock(&data_[i]);
+    }
+
+   private:
+    omp_lock_t* data_ = nullptr;
+    size_t size_ = 0;
+    size_t capacity_ = 0;
+};
+
+} // namespace faiss

data/vendor/faiss/faiss/impl/hnsw/MinimaxHeap.cpp

@@ -0,0 +1,83 @@
+/*
+ * 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/hnsw/MinimaxHeap.h>
+
+#include <faiss/impl/simd_dispatch.h>
+
+namespace faiss {
+
+// Runtime-dispatched pop_min (NONE + AVX2 + AVX512 only).
+constexpr int MINIMAX_HEAP_SIMD_LEVELS = (1 << int(SIMDLevel::NONE)) |
+        (1 << int(SIMDLevel::AVX2)) | (1 << int(SIMDLevel::AVX512));
+
+template <class HC_>
+int MinimaxHeapT<HC_>::pop_min(float* vmin_out) {
+    return with_selected_simd_levels<MINIMAX_HEAP_SIMD_LEVELS>(
+            [&]<SIMDLevel SL>() {
+                return pop_min_tpl<HC_, SL>(this, vmin_out);
+            });
+}
+
+// Primary-template scalar implementation. Used directly when SL==NONE
+template <class HC>
+int pop_min_simd_none(MinimaxHeapT<HC>* heap, float* vmin_out) {
+    int k = heap->k;
+    int* ids = heap->ids.data();
+    float* dis = heap->dis.data();
+    assert(k > 0);
+    // Returns the "best" entry. This is an O(n) operation.
+    int i = k - 1;
+    while (i >= 0) {
+        if (ids[i] != -1) {
+            break;
+        }
+        i--;
+    }
+    if (i == -1) {
+        return -1;
+    }
+    int imin = i;
+    float vmin = dis[i];
+    i--;
+    while (i >= 0) {
+        // HC::cmp(vmin, dis[i]) → "dis[i] is better than vmin".
+        if (ids[i] != -1 && HC::cmp(vmin, dis[i])) {
+            vmin = dis[i];
+            imin = i;
+        }
+        i--;
+    }
+    if (vmin_out) {
+        *vmin_out = vmin;
+    }
+    int ret = ids[imin];
+    ids[imin] = -1;
+    --heap->nvalid;
+    return ret;
+}
+
+// declare for min and max heap at simd level NONE
+template <>
+int pop_min_tpl<CMin<float, int32_t>, SIMDLevel::NONE>(
+        MinimaxHeapT<CMin<float, int32_t>>* heap,
+        float* vmin_out) {
+    return pop_min_simd_none(heap, vmin_out);
+}
+
+template <>
+int pop_min_tpl<CMax<float, int32_t>, SIMDLevel::NONE>(
+        MinimaxHeapT<CMax<float, int32_t>>* heap,
+        float* vmin_out) {
+    return pop_min_simd_none(heap, vmin_out);
+}
+
+// Explicit instantiations of pop_min for the two HC variants
+template int MinimaxHeapT<CMax<float, int32_t>>::pop_min(float*);
+template int MinimaxHeapT<CMin<float, int32_t>>::pop_min(float*);
+
+} // namespace faiss

data/vendor/faiss/faiss/impl/hnsw/MinimaxHeap.h

@@ -0,0 +1,113 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include <cassert>
+#include <cmath>
+#include <cstdint>
+#include <vector>
+
+#include <faiss/utils/Heap.h>
+#include <faiss/utils/ordered_key_value.h>
+#include <faiss/utils/simd_levels.h>
+
+namespace faiss {
+
+/** Heap structure that allows fast access and updates.
+ *
+ * Templated on the comparator HC_ so that the same data structure can
+ * service both distance-style searches (HC_ = CMax<float, int32_t>, smaller
+ * is better) and similarity-style searches (HC_ = CMin<float, int32_t>,
+ * larger is better). For the distance variant the underlying heap is a
+ * max-heap and "pop_min" returns the closest element; for similarity the
+ * underlying heap is a min-heap and "pop_min" returns the most similar
+ * element.
+ */
+template <class HC_ = CMax<float, int32_t>>
+struct MinimaxHeapT {
+    using HC = HC_;
+    using storage_idx_t = int32_t;
+
+    int n;
+    int k;
+    int nvalid;
+
+    std::vector<storage_idx_t> ids;
+    std::vector<float> dis;
+
+    explicit MinimaxHeapT(int n_in)
+            : n(n_in), k(0), nvalid(0), ids(n_in), dis(n_in) {}
+
+    void push(storage_idx_t i, float v) {
+        // Treat NaN distances as the "worst" value so heap ordering is
+        // preserved (insertion is then guaranteed to fall through the
+        // not-better-than-top early-reject branch when the heap is full).
+        if (std::isnan(v)) {
+            v = HC::neutral();
+        }
+        if (k == n) {
+            // top of the heap is the "worst" entry under HC. If the new
+            // value is not strictly better than the worst, drop it.
+            // HC::cmp(top, v) means "v is better than top" for both CMax
+            // (cmp = a > b → top > v → v < top) and CMin (cmp = a < b →
+            // top < v → v > top).
+            if (!HC::cmp(dis[0], v)) {
+                return;
+            }
+            if (ids[0] != -1) {
+                --nvalid;
+            }
+            faiss::heap_pop<HC>(k--, dis.data(), ids.data());
+        }
+        faiss::heap_push<HC>(++k, dis.data(), ids.data(), v, i);
+        ++nvalid;
+    }
+
+    float max() const {
+        return dis[0];
+    }
+
+    int size() const {
+        return nvalid;
+    }
+
+    void clear() {
+        nvalid = k = 0;
+    }
+
+    /// Runtime-dispatched best-element extraction (NONE + AVX2 + AVX512).
+    int pop_min(float* vmin_out = nullptr);
+
+    int count_below(float thresh) {
+        int n_below = 0;
+        for (int i = 0; i < k; i++) {
+            // Count entries that are strictly "better than" thresh.
+            // HC::cmp(thresh, dis[i]) → for CMax: thresh > dis[i]
+            // (i.e., dis[i] < thresh, the historical L2 semantics);
+            // for CMin: thresh < dis[i] (similarity above threshold).
+            if (HC::cmp(thresh, dis[i])) {
+                n_below++;
+            }
+        }
+        return n_below;
+    }
+};
+
+// Default `MinimaxHeap` keeps the historical max-heap semantics (smaller
+// distance is better). The CMin instantiation is used when the owning
+// HNSW has `is_similarity = true`. The alias itself is declared once,
+// alongside the forward declaration in HNSW.h, to avoid duplicate
+// `using` declarations that SWIG treats as redundant.
+
+// Forward declarations of the SIMD specializations. The actual bodies live
+// in the SIMD-specific translation units (avx2.cpp, avx512.cpp) and are
+// resolved at link time.
+template <class HC_, SIMDLevel SL>
+int pop_min_tpl(MinimaxHeapT<HC_>* heap, float* vmin_out);
+
+} // namespace faiss

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

@@ -0,0 +1,150 @@
+/*
+ * 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.
+ */
+
+#ifdef COMPILE_SIMD_AVX2
+
+#include <faiss/impl/hnsw/MinimaxHeap.h>
+
+#include <immintrin.h>
+#include <cassert>
+#include <limits>
+#include <type_traits>
+
+namespace faiss {
+
+namespace {
+
+/// Templated AVX2 implementation of "pop best" for both CMax (returns
+/// the smallest distance) and CMin (returns the largest similarity).
+/// The only differences between the two flavors are: (1) the initial
+/// "worst possible" value, (2) the running-best update comparison
+/// (`_CMP_LT_OS` vs `_CMP_GT_OS`), and (3) the tiebreaker direction.
+template <class HC>
+int pop_best_avx2(MinimaxHeapT<HC>& heap, float* vmin_out) {
+    using storage_idx_t = typename MinimaxHeapT<HC>::storage_idx_t;
+    static_assert(
+            std::is_same<storage_idx_t, int32_t>::value,
+            "This code expects storage_idx_t to be int32_t");
+    assert(heap.k > 0);
+
+    // For CMax (distance) the "best" candidate is the smallest value, so
+    // we initialize the running best to +inf. For CMin (similarity) the
+    // best is the largest value, so we initialize to -inf.
+    constexpr float worst_v = HC::is_max
+            ? std::numeric_limits<float>::infinity()
+            : -std::numeric_limits<float>::infinity();
+
+    int32_t best_idx = -1;
+    float best_dis = worst_v;
+
+    size_t iii = 0;
+
+    __m256i best_indices = _mm256_setr_epi32(-1, -1, -1, -1, -1, -1, -1, -1);
+    __m256 best_distances = _mm256_set1_ps(worst_v);
+    __m256i current_indices = _mm256_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7);
+    __m256i offset = _mm256_set1_epi32(8);
+
+    // Track the rightmost index whose distance equals the running best.
+    // -1 index values are filtered out via m1mask.
+    const size_t k8 = (heap.k / 8) * 8;
+    for (; iii < k8; iii += 8) {
+        __m256i indices =
+                _mm256_loadu_si256((const __m256i*)(heap.ids.data() + iii));
+        __m256 distances = _mm256_loadu_ps(heap.dis.data() + iii);
+
+        // Mask out -1 indices (invalid entries).
+        __m256i m1mask = _mm256_cmpgt_epi32(_mm256_setzero_si256(), indices);
+
+        // dmask is "true where best is already (strictly) better than the
+        // candidate" — entries the candidate should NOT update. For CMax,
+        // best < candidate means we keep best (we want the smallest);
+        // for CMin we keep best when best > candidate (we want the largest).
+        __m256i dmask;
+        if constexpr (HC::is_max) {
+            dmask = _mm256_castps_si256(
+                    _mm256_cmp_ps(best_distances, distances, _CMP_LT_OS));
+        } else {
+            dmask = _mm256_castps_si256(
+                    _mm256_cmp_ps(best_distances, distances, _CMP_GT_OS));
+        }
+        __m256 finalmask = _mm256_castsi256_ps(_mm256_or_si256(m1mask, dmask));
+
+        const __m256i best_indices_new = _mm256_castps_si256(_mm256_blendv_ps(
+                _mm256_castsi256_ps(current_indices),
+                _mm256_castsi256_ps(best_indices),
+                finalmask));
+
+        const __m256 best_distances_new =
+                _mm256_blendv_ps(distances, best_distances, finalmask);
+
+        best_indices = best_indices_new;
+        best_distances = best_distances_new;
+
+        current_indices = _mm256_add_epi32(current_indices, offset);
+    }
+
+    // Vectorizing the horizontal reduction is doable but not practical.
+    int32_t vidx8[8];
+    float vdis8[8];
+    _mm256_storeu_ps(vdis8, best_distances);
+    _mm256_storeu_si256((__m256i*)vidx8, best_indices);
+
+    for (size_t j = 0; j < 8; j++) {
+        const bool strictly_better =
+                HC::is_max ? (best_dis > vdis8[j]) : (best_dis < vdis8[j]);
+        if (strictly_better || (best_dis == vdis8[j] && best_idx < vidx8[j])) {
+            best_idx = vidx8[j];
+            best_dis = vdis8[j];
+        }
+    }
+
+    // Tail (under 8 entries). Vectorizing is doable but not practical.
+    for (; iii < static_cast<size_t>(heap.k); iii++) {
+        if (heap.ids[iii] == -1) {
+            continue;
+        }
+        const bool weakly_better = HC::is_max ? (best_dis >= heap.dis[iii])
+                                              : (best_dis <= heap.dis[iii]);
+        if (weakly_better) {
+            best_dis = heap.dis[iii];
+            best_idx = iii;
+        }
+    }
+
+    if (best_idx == -1) {
+        return -1;
+    }
+
+    if (vmin_out) {
+        *vmin_out = best_dis;
+    }
+    int ret = heap.ids[best_idx];
+    heap.ids[best_idx] = -1;
+    --heap.nvalid;
+    return ret;
+}
+
+} // namespace
+
+// Explicit specializations for AVX2
+template <>
+int pop_min_tpl<CMax<float, int32_t>, SIMDLevel::AVX2>(
+        MinimaxHeapT<CMax<float, int32_t>>* heap,
+        float* vmin_out) {
+    return pop_best_avx2<CMax<float, int32_t>>(*heap, vmin_out);
+}
+
+template <>
+int pop_min_tpl<CMin<float, int32_t>, SIMDLevel::AVX2>(
+        MinimaxHeapT<CMin<float, int32_t>>* heap,
+        float* vmin_out) {
+    return pop_best_avx2<CMin<float, int32_t>>(*heap, vmin_out);
+}
+
+} // namespace faiss
+
+#endif // COMPILE_SIMD_AVX2

data/vendor/faiss/faiss/impl/hnsw/avx512.cpp

@@ -0,0 +1,142 @@
+/*
+ * 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.
+ */
+
+#ifdef COMPILE_SIMD_AVX512
+
+#include <faiss/impl/hnsw/MinimaxHeap.h>
+
+#include <immintrin.h>
+#include <cassert>
+#include <limits>
+#include <type_traits>
+
+namespace faiss {
+
+namespace {
+
+/// Templated AVX512 implementation of "pop best" for both CMax (returns
+/// the smallest distance) and CMin (returns the largest similarity).
+template <class HC>
+int pop_best_avx512(MinimaxHeapT<HC>& heap, float* vmin_out) {
+    using storage_idx_t = typename MinimaxHeapT<HC>::storage_idx_t;
+    static_assert(
+            std::is_same<storage_idx_t, int32_t>::value,
+            "This code expects storage_idx_t to be int32_t");
+    assert(heap.k > 0);
+
+    constexpr float worst_v = HC::is_max
+            ? std::numeric_limits<float>::infinity()
+            : -std::numeric_limits<float>::infinity();
+
+    int32_t best_idx = -1;
+    float best_dis = worst_v;
+
+    __m512i best_indices = _mm512_set1_epi32(-1);
+    __m512 best_distances = _mm512_set1_ps(worst_v);
+    __m512i current_indices = _mm512_setr_epi32(
+            0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15);
+    __m512i offset = _mm512_set1_epi32(16);
+
+    auto best_vs_cand_mask = [](__m512 best_d, __m512 cand_d) -> __mmask16 {
+        // Returns the mask of lanes where the current best is already
+        // (strictly) better than the candidate.
+        if constexpr (HC::is_max) {
+            return _mm512_cmp_ps_mask(best_d, cand_d, _CMP_LT_OS);
+        } else {
+            return _mm512_cmp_ps_mask(best_d, cand_d, _CMP_GT_OS);
+        }
+    };
+
+    const size_t k16 = (heap.k / 16) * 16;
+    for (size_t iii = 0; iii < k16; iii += 16) {
+        __m512i indices =
+                _mm512_loadu_si512((const __m512i*)(heap.ids.data() + iii));
+        __m512 distances = _mm512_loadu_ps(heap.dis.data() + iii);
+
+        __mmask16 m1mask =
+                _mm512_cmpgt_epi32_mask(_mm512_setzero_si512(), indices);
+        __mmask16 dmask = best_vs_cand_mask(best_distances, distances);
+        __mmask16 finalmask = m1mask | dmask;
+
+        const __m512i best_indices_new = _mm512_mask_blend_epi32(
+                finalmask, current_indices, best_indices);
+        const __m512 best_distances_new =
+                _mm512_mask_blend_ps(finalmask, distances, best_distances);
+
+        best_indices = best_indices_new;
+        best_distances = best_distances_new;
+
+        current_indices = _mm512_add_epi32(current_indices, offset);
+    }
+
+    // Leftovers.
+    if (k16 != static_cast<size_t>(heap.k)) {
+        const __mmask16 kmask = (1 << (heap.k - k16)) - 1;
+
+        __m512i indices = _mm512_mask_loadu_epi32(
+                _mm512_set1_epi32(-1), kmask, heap.ids.data() + k16);
+        __m512 distances = _mm512_maskz_loadu_ps(kmask, heap.dis.data() + k16);
+
+        __mmask16 m1mask =
+                _mm512_cmpgt_epi32_mask(_mm512_setzero_si512(), indices);
+        __mmask16 dmask = best_vs_cand_mask(best_distances, distances);
+        __mmask16 finalmask = m1mask | dmask;
+
+        const __m512i best_indices_new = _mm512_mask_blend_epi32(
+                finalmask, current_indices, best_indices);
+        const __m512 best_distances_new =
+                _mm512_mask_blend_ps(finalmask, distances, best_distances);
+
+        best_indices = best_indices_new;
+        best_distances = best_distances_new;
+    }
+
+    // Horizontal best: min for CMax (distance), max for CMin (similarity).
+    if constexpr (HC::is_max) {
+        best_dis = _mm512_reduce_min_ps(best_distances);
+    } else {
+        best_dis = _mm512_reduce_max_ps(best_distances);
+    }
+    // Tiebreak by picking the rightmost (largest) index among lanes
+    // matching the best distance, matching the original behavior.
+    __mmask16 best_lane_mask =
+            _mm512_cmpeq_ps_mask(best_distances, _mm512_set1_ps(best_dis));
+    best_idx = _mm512_mask_reduce_max_epi32(best_lane_mask, best_indices);
+
+    if (best_idx == -1) {
+        return -1;
+    }
+
+    if (vmin_out) {
+        *vmin_out = best_dis;
+    }
+    int ret = heap.ids[best_idx];
+    heap.ids[best_idx] = -1;
+    --heap.nvalid;
+    return ret;
+}
+
+} // namespace
+
+// Explicit specializations for AVX512
+template <>
+int pop_min_tpl<CMax<float, int32_t>, SIMDLevel::AVX512>(
+        MinimaxHeapT<CMax<float, int32_t>>* heap,
+        float* vmin_out) {
+    return pop_best_avx512<CMax<float, int32_t>>(*heap, vmin_out);
+}
+
+template <>
+int pop_min_tpl<CMin<float, int32_t>, SIMDLevel::AVX512>(
+        MinimaxHeapT<CMin<float, int32_t>>* heap,
+        float* vmin_out) {
+    return pop_best_avx512<CMin<float, int32_t>>(*heap, vmin_out);
+}
+
+} // namespace faiss
+
+#endif // COMPILE_SIMD_AVX512