HNSW SIMD dispatch (#5015)

Summary:
Pull Request resolved: #5015

Extract `HNSW::MinimaxHeap` into a standalone `faiss::MinimaxHeap` struct in a new `impl/hnsw/` subdirectory, and convert `pop_min` to use FAISS SIMD dynamic dispatch.

### Changes

**New files:**
- `impl/hnsw/MinimaxHeap.h` — Standalone `faiss::MinimaxHeap` struct (no longer nested in `HNSW`). Trivial methods (`max`, `size`, `clear`) inlined. Declares `template<SIMDLevel> pop_min_tpl` and runtime-dispatched `pop_min`.
- `impl/hnsw/MinimaxHeap.cpp` — `push`, `count_below`, scalar `pop_min_tpl<NONE>`, and `pop_min` dispatch wrapper using `with_selected_simd_levels`.
- `impl/hnsw/avx2.cpp` — `pop_min_tpl<AVX2>` specialization, guarded by `COMPILE_SIMD_AVX2`.
- `impl/hnsw/avx512.cpp` — `pop_min_tpl<AVX512>` specialization, guarded by `COMPILE_SIMD_AVX512`.
- `impl/hnsw/.nobuck` — Prevents autodeps from generating a BUCK file (sources are part of top-level faiss target).

**Modified files:**
- `impl/HNSW.h` — Includes new header, removed nested struct, free functions use `MinimaxHeap&`.
- `impl/HNSW.cpp` — Removed all `MinimaxHeap` implementations.
- `IndexHNSW.cpp` — Removed stale `using MinimaxHeap = HNSW::MinimaxHeap` alias.
- `tests/test_hnsw.cpp` — Updated `faiss::HNSW::MinimaxHeap` → `faiss::MinimaxHeap`.
- `CMakeLists.txt` — Added new source/header files and SIMD file registrations.
- `xplat.bzl` — Added new source/header files and SIMD_FILES entries.

Reviewed By: algoriddle

Differential Revision: D98920170

fbshipit-source-id: 7c4d04e7eb6647b359efb2fa2f86c9f2a084d189

Author: mdouze

faiss/CMakeLists.txt

@@ -10,6 +10,7 @@
 # =============================================================================
 set(FAISS_SIMD_AVX2_SRC
   impl/fast_scan/impl-avx2.cpp
+  impl/hnsw/avx2.cpp
   impl/pq_code_distance/pq_code_distance-avx2.cpp
   impl/scalar_quantizer/sq-avx2.cpp
   impl/approx_topk/avx2.cpp
@@ -19,6 +20,7 @@ set(FAISS_SIMD_AVX2_SRC
 )
 set(FAISS_SIMD_AVX512_SRC
   impl/fast_scan/impl-avx512.cpp
+  impl/hnsw/avx512.cpp
   impl/pq_code_distance/pq_code_distance-avx512.cpp
   impl/scalar_quantizer/sq-avx512.cpp
   utils/simd_impl/distances_avx512.cpp
@@ -108,6 +110,7 @@ set(FAISS_SRC
   impl/IDSelector.cpp
   impl/FaissException.cpp
   impl/HNSW.cpp
+  impl/hnsw/MinimaxHeap.cpp
   impl/NSG.cpp
   impl/PolysemousTraining.cpp
   impl/ProductQuantizer.cpp
@@ -232,6 +235,7 @@ set(FAISS_HEADERS
   impl/FaissAssert.h
   impl/FaissException.h
   impl/HNSW.h
+  impl/hnsw/MinimaxHeap.h
   impl/LocalSearchQuantizer.h
   impl/ProductAdditiveQuantizer.h
   impl/fast_scan/LookupTableScaler.h

faiss/IndexHNSW.cpp

@@ -33,7 +33,6 @@
 
 namespace faiss {
 
-using MinimaxHeap = HNSW::MinimaxHeap;
 using storage_idx_t = HNSW::storage_idx_t;
 using NodeDistFarther = HNSW::NodeDistFarther;
 

faiss/impl/HNSW.cpp

@@ -17,13 +17,6 @@
 #include <faiss/impl/ResultHandler.h>
 #include <faiss/impl/VisitedTable.h>
 
-#ifdef __AVX2__
-#include <immintrin.h>
-
-#include <limits>
-#include <type_traits>
-#endif
-
 namespace faiss {
 
 /**************************************************************
@@ -597,7 +590,6 @@ void HNSW::add_with_locks(
  * Searching
  **************************************************************/
 
-using MinimaxHeap = HNSW::MinimaxHeap;
 using Node = HNSW::Node;
 using C = HNSW::C;
 
@@ -1190,7 +1182,6 @@ HNSWStats greedy_update_nearest(
 }
 
 namespace {
-using MinimaxHeap = HNSW::MinimaxHeap;
 using Node = HNSW::Node;
 using C = HNSW::C;
 
@@ -1388,257 +1379,4 @@ void HNSW::permute_entries(const idx_t* map) {
     neighbors = std::move(new_neighbors);
 }
 
-/**************************************************************
- * MinimaxHeap
- **************************************************************/
-
-void HNSW::MinimaxHeap::push(storage_idx_t i, float v) {
-    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;
-}
-
-float HNSW::MinimaxHeap::max() const {
-    return dis[0];
-}
-
-int HNSW::MinimaxHeap::size() const {
-    return nvalid;
-}
-
-void HNSW::MinimaxHeap::clear() {
-    nvalid = k = 0;
-}
-
-#ifdef __AVX512F__
-
-int HNSW::MinimaxHeap::pop_min(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;
-}
-
-#elif __AVX2__
-
-int HNSW::MinimaxHeap::pop_min(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 non-AVX2 branch.
-
-    // 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;
-}
-
-#else
-
-// baseline non-vectorized version
-int HNSW::MinimaxHeap::pop_min(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;
-}
-#endif
-
-int HNSW::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

faiss/impl/HNSW.h

@@ -16,6 +16,7 @@
 #include <faiss/Index.h>
 #include <faiss/impl/DistanceComputer.h>
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/hnsw/MinimaxHeap.h>
 #include <faiss/impl/maybe_owned_vector.h>
 #include <faiss/impl/platform_macros.h>
 #include <faiss/utils/Heap.h>
@@ -64,33 +65,6 @@ struct HNSW {
 
     typedef std::pair<float, storage_idx_t> Node;
 
-    /** Heap structure that allows fast access and updates.
-     */
-    struct MinimaxHeap {
-        int n;
-        int k;
-        int nvalid;
-
-        std::vector<storage_idx_t> ids;
-        std::vector<float> dis;
-        typedef faiss::CMax<float, storage_idx_t> HC;
-
-        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;
-
-        int size() const;
-
-        void clear();
-
-        int pop_min(float* vmin_out = nullptr);
-
-        int count_below(float thresh);
-    };
-
     /// to sort pairs of (id, distance) from nearest to farthest or the reverse
     struct NodeDistCloser {
         float d;
@@ -280,7 +254,7 @@ int search_from_candidates(
         const HNSW& hnsw,
         DistanceComputer& qdis,
         ResultHandler& res,
-        HNSW::MinimaxHeap& candidates,
+        MinimaxHeap& candidates,
         VisitedTable& vt,
         HNSWStats& stats,
         int level,
@@ -296,7 +270,7 @@ int search_from_candidates_panorama(
         const IndexHNSW* index,
         DistanceComputer& qdis,
         ResultHandler& res,
-        HNSW::MinimaxHeap& candidates,
+        MinimaxHeap& candidates,
         VisitedTable& vt,
         HNSWStats& stats,
         int level,

faiss/impl/hnsw/.nobuck

@@ -0,0 +1,2 @@
+# This directory's sources are compiled as part of the top-level faiss target.
+# No separate BUCK file is needed here.