faiss 0.6.1 → 0.6.3

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

@@ -16,89 +16,135 @@
 
 namespace faiss {
 
-template <>
-int MinimaxHeap::pop_min_tpl<SIMDLevel::AVX2>(float* vmin_out) {
-    assert(k > 0);
+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 min_idx = -1;
-    float min_dis = std::numeric_limits<float>::infinity();
+    int32_t best_idx = -1;
+    float best_dis = worst_v;
 
     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 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);
 
-    // 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;
+    // 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*)(ids.data() + iii));
-        __m256 distances = _mm256_loadu_ps(dis.data() + iii);
+                _mm256_loadu_si256((const __m256i*)(heap.ids.data() + iii));
+        __m256 distances = _mm256_loadu_ps(heap.dis.data() + iii);
 
-        // This mask filters out -1 values among indices.
+        // Mask out -1 indices (invalid entries).
         __m256i m1mask = _mm256_cmpgt_epi32(_mm256_setzero_si256(), indices);
 
-        __m256i dmask = _mm256_castps_si256(
-                _mm256_cmp_ps(min_distances, distances, _CMP_LT_OS));
+        // 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 min_indices_new = _mm256_castps_si256(_mm256_blendv_ps(
+        const __m256i best_indices_new = _mm256_castps_si256(_mm256_blendv_ps(
                 _mm256_castsi256_ps(current_indices),
-                _mm256_castsi256_ps(min_indices),
+                _mm256_castsi256_ps(best_indices),
                 finalmask));
 
-        const __m256 min_distances_new =
-                _mm256_blendv_ps(distances, min_distances, finalmask);
+        const __m256 best_distances_new =
+                _mm256_blendv_ps(distances, best_distances, finalmask);
 
-        min_indices = min_indices_new;
-        min_distances = min_distances_new;
+        best_indices = best_indices_new;
+        best_distances = best_distances_new;
 
         current_indices = _mm256_add_epi32(current_indices, offset);
     }
 
-    // Vectorizing is doable, but is not practical
+    // Vectorizing the horizontal reduction is doable but not practical.
     int32_t vidx8[8];
     float vdis8[8];
-    _mm256_storeu_ps(vdis8, min_distances);
-    _mm256_storeu_si256((__m256i*)vidx8, min_indices);
+    _mm256_storeu_ps(vdis8, best_distances);
+    _mm256_storeu_si256((__m256i*)vidx8, best_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];
+        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];
         }
     }
 
-    // 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;
+    // 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 (min_idx == -1) {
+    if (best_idx == -1) {
         return -1;
     }
 
     if (vmin_out) {
-        *vmin_out = min_dis;
+        *vmin_out = best_dis;
     }
-    int ret = ids[min_idx];
-    ids[min_idx] = -1;
-    --nvalid;
+    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

@@ -16,96 +16,127 @@
 
 namespace faiss {
 
-template <>
-int MinimaxHeap::pop_min_tpl<SIMDLevel::AVX512>(float* vmin_out) {
-    assert(k > 0);
+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);
 
-    int32_t min_idx = -1;
-    float min_dis = std::numeric_limits<float>::infinity();
+    constexpr float worst_v = HC::is_max
+            ? std::numeric_limits<float>::infinity()
+            : -std::numeric_limits<float>::infinity();
 
-    __m512i min_indices = _mm512_set1_epi32(-1);
-    __m512 min_distances =
-            _mm512_set1_ps(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);
 
-    // The following loop tracks the rightmost index with the min distance.
-    // -1 index values are ignored.
-    const size_t k16 = (k / 16) * 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*)(ids.data() + iii));
-        __m512 distances = _mm512_loadu_ps(dis.data() + iii);
+                _mm512_loadu_si512((const __m512i*)(heap.ids.data() + iii));
+        __m512 distances = _mm512_loadu_ps(heap.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 dmask = best_vs_cand_mask(best_distances, distances);
         __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);
+        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);
 
-        min_indices = min_indices_new;
-        min_distances = min_distances_new;
+        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>(k)) {
-        const __mmask16 kmask = (1 << (k - k16)) - 1;
+    // 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, ids.data() + k16);
-        __m512 distances = _mm512_maskz_loadu_ps(kmask, dis.data() + k16);
+                _mm512_set1_epi32(-1), kmask, heap.ids.data() + k16);
+        __m512 distances = _mm512_maskz_loadu_ps(kmask, heap.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 dmask = best_vs_cand_mask(best_distances, distances);
         __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);
+        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);
 
-        min_indices = min_indices_new;
-        min_distances = min_distances_new;
+        best_indices = best_indices_new;
+        best_distances = best_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);
+    // 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 (min_idx == -1) {
+    if (best_idx == -1) {
         return -1;
     }
 
     if (vmin_out) {
-        *vmin_out = min_dis;
+        *vmin_out = best_dis;
     }
-    int ret = ids[min_idx];
-    ids[min_idx] = -1;
-    --nvalid;
+    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