faiss 0.6.0 → 0.6.1

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,91 @@
+/*
+ * 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 <cmath>
+
+#include <faiss/impl/hnsw/MinimaxHeap.h>
+
+#include <cassert>
+
+#include <faiss/impl/simd_dispatch.h>
+
+namespace faiss {
+
+void MinimaxHeap::push(storage_idx_t i, float v) {
+    // Treat NaN distances as infinitely far away so heap ordering is preserved.
+    if (std::isnan(v)) {
+        v = HC::neutral();
+    }
+    if (k == n) {
+        if (v >= dis[0]) {
+            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;
+}
+
+// Scalar (NONE) specialization of pop_min_tpl
+template <>
+int MinimaxHeap::pop_min_tpl<SIMDLevel::NONE>(float* vmin_out) {
+    assert(k > 0);
+    // returns min. 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) {
+        if (ids[i] != -1 && dis[i] < vmin) {
+            vmin = dis[i];
+            imin = i;
+        }
+        i--;
+    }
+    if (vmin_out) {
+        *vmin_out = vmin;
+    }
+    int ret = ids[imin];
+    ids[imin] = -1;
+    --nvalid;
+
+    return ret;
+}
+
+// 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));
+
+int MinimaxHeap::pop_min(float* vmin_out) {
+    return with_selected_simd_levels<MINIMAX_HEAP_SIMD_LEVELS>(
+            [&]<SIMDLevel SL>() { return pop_min_tpl<SL>(vmin_out); });
+}
+
+int MinimaxHeap::count_below(float thresh) {
+    int n_below = 0;
+    for (int i = 0; i < k; i++) {
+        if (dis[i] < thresh) {
+            n_below++;
+        }
+    }
+
+    return n_below;
+}
+
+} // namespace faiss

data/vendor/faiss/faiss/impl/hnsw/MinimaxHeap.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 <cstdint>
+#include <vector>
+
+#include <faiss/utils/Heap.h>
+#include <faiss/utils/simd_levels.h>
+
+namespace faiss {
+
+/** Heap structure that allows fast access and updates.
+ *
+ * Supports both max-heap operations (via the underlying CMax heap)
+ * and efficient min extraction via linear scan (with optional SIMD
+ * acceleration).
+ */
+struct MinimaxHeap {
+    using storage_idx_t = int32_t;
+
+    int n;
+    int k;
+    int nvalid;
+
+    std::vector<storage_idx_t> ids;
+    std::vector<float> dis;
+    using HC = faiss::CMax<float, storage_idx_t>;
+
+    explicit MinimaxHeap(int n_in)
+            : n(n_in), k(0), nvalid(0), ids(n_in), dis(n_in) {}
+
+    void push(storage_idx_t i, float v);
+
+    float max() const {
+        return dis[0];
+    }
+
+    int size() const {
+        return nvalid;
+    }
+
+    void clear() {
+        nvalid = k = 0;
+    }
+
+    /// SIMD-templated pop_min implementation.
+    /// Specializations exist for NONE, AVX2, and AVX512.
+    template <SIMDLevel SL>
+    int pop_min_tpl(float* vmin_out = nullptr);
+
+    /// Runtime-dispatched pop_min (calls pop_min_tpl with best available
+    /// SIMD level).
+    int pop_min(float* vmin_out = nullptr);
+
+    int count_below(float thresh);
+};
+
+} // namespace faiss

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

@@ -0,0 +1,104 @@
+/*
+ * 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 {
+
+template <>
+int MinimaxHeap::pop_min_tpl<SIMDLevel::AVX2>(float* vmin_out) {
+    assert(k > 0);
+    static_assert(
+            std::is_same<storage_idx_t, int32_t>::value,
+            "This code expects storage_idx_t to be int32_t");
+
+    int32_t min_idx = -1;
+    float min_dis = std::numeric_limits<float>::infinity();
+
+    size_t iii = 0;
+
+    __m256i min_indices = _mm256_setr_epi32(-1, -1, -1, -1, -1, -1, -1, -1);
+    __m256 min_distances =
+            _mm256_set1_ps(std::numeric_limits<float>::infinity());
+    __m256i current_indices = _mm256_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7);
+    __m256i offset = _mm256_set1_epi32(8);
+
+    // The baseline version is available in the NONE specialization.
+
+    // The following loop tracks the rightmost index with the min distance.
+    // -1 index values are ignored.
+    const size_t k8 = (k / 8) * 8;
+    for (; iii < k8; iii += 8) {
+        __m256i indices =
+                _mm256_loadu_si256((const __m256i*)(ids.data() + iii));
+        __m256 distances = _mm256_loadu_ps(dis.data() + iii);
+
+        // This mask filters out -1 values among indices.
+        __m256i m1mask = _mm256_cmpgt_epi32(_mm256_setzero_si256(), indices);
+
+        __m256i dmask = _mm256_castps_si256(
+                _mm256_cmp_ps(min_distances, distances, _CMP_LT_OS));
+        __m256 finalmask = _mm256_castsi256_ps(_mm256_or_si256(m1mask, dmask));
+
+        const __m256i min_indices_new = _mm256_castps_si256(_mm256_blendv_ps(
+                _mm256_castsi256_ps(current_indices),
+                _mm256_castsi256_ps(min_indices),
+                finalmask));
+
+        const __m256 min_distances_new =
+                _mm256_blendv_ps(distances, min_distances, finalmask);
+
+        min_indices = min_indices_new;
+        min_distances = min_distances_new;
+
+        current_indices = _mm256_add_epi32(current_indices, offset);
+    }
+
+    // Vectorizing is doable, but is not practical
+    int32_t vidx8[8];
+    float vdis8[8];
+    _mm256_storeu_ps(vdis8, min_distances);
+    _mm256_storeu_si256((__m256i*)vidx8, min_indices);
+
+    for (size_t j = 0; j < 8; j++) {
+        if (min_dis > vdis8[j] || (min_dis == vdis8[j] && min_idx < vidx8[j])) {
+            min_idx = vidx8[j];
+            min_dis = vdis8[j];
+        }
+    }
+
+    // process last values. Vectorizing is doable, but is not practical
+    for (; iii < static_cast<size_t>(k); iii++) {
+        if (ids[iii] != -1 && dis[iii] <= min_dis) {
+            min_dis = dis[iii];
+            min_idx = iii;
+        }
+    }
+
+    if (min_idx == -1) {
+        return -1;
+    }
+
+    if (vmin_out) {
+        *vmin_out = min_dis;
+    }
+    int ret = ids[min_idx];
+    ids[min_idx] = -1;
+    --nvalid;
+    return ret;
+}
+
+} // namespace faiss
+
+#endif // COMPILE_SIMD_AVX2

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

@@ -0,0 +1,111 @@
+/*
+ * 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 {
+
+template <>
+int MinimaxHeap::pop_min_tpl<SIMDLevel::AVX512>(float* vmin_out) {
+    assert(k > 0);
+    static_assert(
+            std::is_same<storage_idx_t, int32_t>::value,
+            "This code expects storage_idx_t to be int32_t");
+
+    int32_t min_idx = -1;
+    float min_dis = std::numeric_limits<float>::infinity();
+
+    __m512i min_indices = _mm512_set1_epi32(-1);
+    __m512 min_distances =
+            _mm512_set1_ps(std::numeric_limits<float>::infinity());
+    __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);
+
+    // The following loop tracks the rightmost index with the min distance.
+    // -1 index values are ignored.
+    const size_t k16 = (k / 16) * 16;
+    for (size_t iii = 0; iii < k16; iii += 16) {
+        __m512i indices =
+                _mm512_loadu_si512((const __m512i*)(ids.data() + iii));
+        __m512 distances = _mm512_loadu_ps(dis.data() + iii);
+
+        // This mask filters out -1 values among indices.
+        __mmask16 m1mask =
+                _mm512_cmpgt_epi32_mask(_mm512_setzero_si512(), indices);
+
+        __mmask16 dmask =
+                _mm512_cmp_ps_mask(min_distances, distances, _CMP_LT_OS);
+        __mmask16 finalmask = m1mask | dmask;
+
+        const __m512i min_indices_new = _mm512_mask_blend_epi32(
+                finalmask, current_indices, min_indices);
+        const __m512 min_distances_new =
+                _mm512_mask_blend_ps(finalmask, distances, min_distances);
+
+        min_indices = min_indices_new;
+        min_distances = min_distances_new;
+
+        current_indices = _mm512_add_epi32(current_indices, offset);
+    }
+
+    // leftovers
+    if (k16 != static_cast<size_t>(k)) {
+        const __mmask16 kmask = (1 << (k - k16)) - 1;
+
+        __m512i indices = _mm512_mask_loadu_epi32(
+                _mm512_set1_epi32(-1), kmask, ids.data() + k16);
+        __m512 distances = _mm512_maskz_loadu_ps(kmask, dis.data() + k16);
+
+        // This mask filters out -1 values among indices.
+        __mmask16 m1mask =
+                _mm512_cmpgt_epi32_mask(_mm512_setzero_si512(), indices);
+
+        __mmask16 dmask =
+                _mm512_cmp_ps_mask(min_distances, distances, _CMP_LT_OS);
+        __mmask16 finalmask = m1mask | dmask;
+
+        const __m512i min_indices_new = _mm512_mask_blend_epi32(
+                finalmask, current_indices, min_indices);
+        const __m512 min_distances_new =
+                _mm512_mask_blend_ps(finalmask, distances, min_distances);
+
+        min_indices = min_indices_new;
+        min_distances = min_distances_new;
+    }
+
+    // grab min distance
+    min_dis = _mm512_reduce_min_ps(min_distances);
+    // blend
+    __mmask16 mindmask =
+            _mm512_cmpeq_ps_mask(min_distances, _mm512_set1_ps(min_dis));
+    // pick the max one
+    min_idx = _mm512_mask_reduce_max_epi32(mindmask, min_indices);
+
+    if (min_idx == -1) {
+        return -1;
+    }
+
+    if (vmin_out) {
+        *vmin_out = min_dis;
+    }
+    int ret = ids[min_idx];
+    ids[min_idx] = -1;
+    --nvalid;
+    return ret;
+}
+
+} // namespace faiss
+
+#endif // COMPILE_SIMD_AVX512