Avoid search stats contention during concurrent search (#5663)

Summary:
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5663

Follow-up to D121097648 (public PR: https://github.com/facebookresearch/faiss/pull/5657).

Related public PR: https://github.com/facebookresearch/faiss/pull/4910

Use relaxed atomic additions for HNSW, RaBitQ, and IVFPQ stats aggregation. Builds with `std::atomic_ref` support use `fetch_add(..., std::memory_order_relaxed)`; older standard libraries fall back to `__atomic_fetch_add(..., __ATOMIC_RELAXED)`. This removes the global OpenMP critical sections that serialize concurrent searches while preserving exact counter updates.

Benchmark design: mode/opt `bench_hnsw_stats` binary, one shared HNSW index, `N` external pthreads, one OpenMP thread per search (each worker calls `omp_set_num_threads(1)` once before its timed loop), five repetitions per point.

Performance:

| External pthreads | Critical M/s | Relaxed atomic M/s | Speedup |
|---:|---:|---:|---:|
| 1 | 2.228 | 2.503 | 1.12x |
| 2 | 3.716 | 3.853 | 1.04x |
| 4 | 5.662 | 7.357 | 1.30x |
| 8 | 6.384 | 12.843 | 2.01x |
| 16 | 2.697 | 15.804 | 5.86x |
| 32 | 0.951 | 18.910 | 19.88x |
| 64 | 0.117 | 21.866 | 186.89x |

At 64 pthreads, the feature-detected `atomic_ref` path measured 21.918 M/s mean over 10 repetitions (21.907 M/s median, 2.59% CV). The `__atomic_fetch_add` fallback measured 22.563 M/s mean (22.511 M/s median, 3.24% CV).

The 64-thread profile attributes 97.38% of cycles in the critical-section version to lock acquisition and waiting. With relaxed atomic aggregation, 74.14% is in the stats combine path instead. The remaining non-linear scaling, 21.866 M/s at 64 threads versus the 2.503 M/s single-thread baseline, comes from cache-line contention on the global counters and shared-index memory-bandwidth and cache pressure.

Reviewed By: mnorris11

Differential Revision: D121368933

fbshipit-source-id: 8dd26313e55aa73c0cdd379c7648efacd4f5f1d8

benchs/bench_hnsw_stats.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.
+ */
+
+#include <array>
+#include <cstdint>
+#include <vector>
+
+#include <benchmark/benchmark.h>
+#include <faiss/IndexBinaryHNSW.h>
+#include <omp.h>
+
+namespace faiss {
+namespace {
+
+constexpr int kDimension = 8;
+constexpr idx_t kDatabaseSize = 1;
+
+struct BenchmarkIndex {
+    BenchmarkIndex() : index(kDimension, 2) {
+        std::vector<uint8_t> database(kDatabaseSize * kDimension / 8);
+        for (size_t i = 0; i < database.size(); ++i) {
+            database[i] = static_cast<uint8_t>(i * 17 + 3);
+        }
+
+        index.hnsw.efConstruction = 2;
+        index.add(kDatabaseSize, database.data());
+        index.hnsw.efSearch = 1;
+    }
+
+    IndexBinaryHNSW index;
+};
+
+const IndexBinaryHNSW& get_index() {
+    static const BenchmarkIndex benchmark_index;
+    return benchmark_index.index;
+}
+
+void bench_binary_hnsw_single_query(benchmark::State& state) {
+    omp_set_num_threads(1);
+    constexpr std::array<uint8_t, kDimension / 8> query{};
+    int32_t distance;
+    idx_t label;
+
+    for (auto _ : state) {
+        get_index().search(1, query.data(), 1, &distance, &label);
+        benchmark::DoNotOptimize(distance);
+        benchmark::DoNotOptimize(label);
+    }
+    state.SetItemsProcessed(state.iterations());
+}
+
+BENCHMARK(bench_binary_hnsw_single_query)->ThreadRange(1, 64)->UseRealTime();
+
+} // namespace
+} // namespace faiss
+
+int main(int argc, char** argv) {
+    omp_set_num_threads(1);
+    faiss::get_index();
+    benchmark::Initialize(&argc, argv);
+    benchmark::RunSpecifiedBenchmarks();
+    benchmark::Shutdown();
+    return 0;
+}

faiss/IndexBinaryHNSW.cpp

@@ -121,10 +121,7 @@ void IndexBinaryHNSW::search(
         }
     }
 
-#pragma omp critical
-    {
-        hnsw_stats.combine({n1, n2, ndis, nhops});
-    }
+    hnsw_stats.combine_atomic({n1, n2, ndis, nhops});
 
 #pragma omp parallel for
     for (idx_t i = 0; i < n * k; ++i) {
@@ -273,10 +270,7 @@ void IndexBinaryHNSWCagra::search(
 
                 res.end();
             }
-#pragma omp critical
-            {
-                hnsw_stats.combine(search_stats);
-            }
+            hnsw_stats.combine_atomic(search_stats);
         }
 
 #pragma omp parallel for

faiss/IndexHNSW.cpp

@@ -32,6 +32,7 @@
 #include <faiss/impl/VisitedTable.h>
 #include <faiss/impl/hnsw/LockVector.h>
 #include <faiss/impl/hnsw/MinimaxHeap.h>
+#include <faiss/impl/platform_macros.h>
 #include <faiss/utils/random.h>
 #include <faiss/utils/sorting.h>
 
@@ -42,6 +43,13 @@ using NodeDistFarther = HNSW::NodeDistFarther;
 
 HNSWStats hnsw_stats;
 
+void HNSWStats::combine_atomic(const HNSWStats& other) {
+    detail::atomic_fetch_add_relaxed(n1, other.n1);
+    detail::atomic_fetch_add_relaxed(n2, other.n2);
+    detail::atomic_fetch_add_relaxed(ndis, other.ndis);
+    detail::atomic_fetch_add_relaxed(nhops, other.nhops);
+}
+
 /**************************************************************
  * add / search blocks of descriptors
  **************************************************************/
@@ -487,11 +495,8 @@ void hnsw_search(
         InterruptCallback::check();
     }
 
-#pragma omp critical
-    {
-        hnsw_stats.combine({n1, n2, ndis, nhops});
-        rabitq_stats.add({n_rabitq_1bit, n_rabitq_refine});
-    }
+    hnsw_stats.combine_atomic({n1, n2, ndis, nhops});
+    rabitq_stats.add_atomic({n_rabitq_1bit, n_rabitq_refine});
 }
 
 } // anonymous namespace
@@ -702,11 +707,8 @@ void IndexHNSW::search_level_0(
                     omp_capture_exception(ex, [&] { interrupt = true; });
                 }
             }
-#pragma omp critical
-            {
-                hnsw_stats.combine(search_stats);
-                rabitq_stats.add(rq_search_stats);
-            }
+            hnsw_stats.combine_atomic(search_stats);
+            rabitq_stats.add_atomic(rq_search_stats);
         }
         omp_rethrow_if_exception(ex);
     };
@@ -1233,10 +1235,7 @@ void IndexHNSW2Level::search(
         }
         omp_rethrow_if_exception(ex);
 
-#pragma omp critical
-        {
-            hnsw_stats.combine({n1, n2, ndis, nhops});
-        }
+        hnsw_stats.combine_atomic({n1, n2, ndis, nhops});
     }
 }
 
@@ -1438,10 +1437,7 @@ void IndexHNSWCagra::range_search(
         result->do_allocation();
         pres.copy_result();
 
-#pragma omp critical
-        {
-            hnsw_stats.combine({n1, n2, ndis, nhops});
-        }
+        hnsw_stats.combine_atomic({n1, n2, ndis, nhops});
     };
 
     if (is_similarity_metric(metric_type)) {

faiss/impl/HNSW.h

@@ -315,6 +315,8 @@ struct HNSWStats {
         ndis += other.ndis;
         nhops += other.nhops;
     }
+
+    void combine_atomic(const HNSWStats& other);
 };
 
 // global var that collects them all

faiss/impl/RaBitQuantizer.cpp

@@ -12,6 +12,7 @@
 #include <faiss/impl/RaBitQUtils.h>
 #include <faiss/impl/RaBitQuantizerMultiBit.h>
 #include <faiss/impl/ResultHandler.h>
+#include <faiss/impl/platform_macros.h>
 #include <faiss/impl/simd_dispatch.h>
 #include <faiss/invlists/DirectMap.h>
 #include <faiss/utils/distances.h>
@@ -27,6 +28,11 @@ namespace faiss {
 
 RaBitQStats rabitq_stats;
 
+void RaBitQStats::add_atomic(const RaBitQStats& other) {
+    detail::atomic_fetch_add_relaxed(n_1bit, other.n_1bit);
+    detail::atomic_fetch_add_relaxed(n_refine, other.n_refine);
+}
+
 // Import shared utilities from RaBitQUtils
 using rabitq_utils::ExtraBitsFactors;
 using rabitq_utils::QueryFactorsData;

faiss/impl/RaBitQuantizer.h

@@ -40,6 +40,8 @@ struct RaBitQStats {
         n_refine += other.n_refine;
     }
 
+    void add_atomic(const RaBitQStats& other);
+
     double refine_ratio() const {
         return n_1bit ? double(n_refine) / double(n_1bit) : 0.0;
     }

faiss/impl/platform_macros.h

@@ -8,6 +8,7 @@
 #pragma once
 
 // basic int types and size_t
+#include <atomic>
 #include <cstdint>
 #include <cstdio>
 
@@ -146,6 +147,21 @@ inline int __builtin_clzll(uint64_t x) {
 #define FAISS_DEPRECATED(msg)
 #endif // GCC or Clang
 
+namespace faiss {
+namespace detail {
+
+template <typename T>
+inline void atomic_fetch_add_relaxed(T& target, T value) {
+#if defined(__cpp_lib_atomic_ref) && __cpp_lib_atomic_ref >= 201806L
+    std::atomic_ref<T>(target).fetch_add(value, std::memory_order_relaxed);
+#else
+    __atomic_fetch_add(&target, value, __ATOMIC_RELAXED);
+#endif
+}
+
+} // namespace detail
+} // namespace faiss
+
 // Localized enablement of imprecise floating point operations
 // You need to use all 3 macros to cover all compilers.
 #if defined(_MSC_VER)

faiss/impl/pq_code_distance/IVFPQScanner_impl.h

@@ -14,6 +14,7 @@
 #include <faiss/IndexIVFPQ.h>
 #include <faiss/impl/IDSelector.h>
 #include <faiss/impl/ResultHandler.h>
+#include <faiss/impl/platform_macros.h>
 #include <faiss/impl/pq_code_distance/IVFPQ_QueryTables.h>
 #include <faiss/impl/pq_code_distance/pq_code_distance-inl.h>
 #include <faiss/impl/simd_dispatch.h>
@@ -411,10 +412,8 @@ struct IVFPQScannerT : QueryTables {
             }
         }
 
-#pragma omp critical
-        {
-            indexIVFPQ_stats.n_hamming_pass += n_hamming_pass;
-        }
+        detail::atomic_fetch_add_relaxed(
+                indexIVFPQ_stats.n_hamming_pass, n_hamming_pass);
     }
 
     template <class SearchResultType>

tests/test_hnsw.cpp

@@ -11,9 +11,11 @@
 #include <cstddef>
 #include <limits>
 #include <random>
+#include <thread>
 #include <unordered_set>
 #include <vector>
 
+#include <faiss/IndexBinaryHNSW.h>
 #include <faiss/IndexFlat.h>
 #include <faiss/IndexHNSW.h>
 #include <faiss/impl/FaissAssert.h>
@@ -22,6 +24,36 @@
 #include <faiss/impl/VisitedTable.h>
 #include <faiss/impl/hnsw/MinimaxHeap.h>
 #include <faiss/utils/random.h>
+#include <omp.h>
+
+TEST(HNSWStats, ConcurrentBinarySearchAggregatesExactly) {
+    constexpr size_t num_threads = 16;
+    constexpr size_t searches_per_thread = 100;
+
+    faiss::IndexBinaryHNSW index(8, 2);
+    const uint8_t database = 0;
+    index.add(1, &database);
+    index.hnsw.efSearch = 1;
+    faiss::hnsw_stats.reset();
+
+    std::vector<std::thread> threads;
+    threads.reserve(num_threads);
+    for (size_t i = 0; i < num_threads; ++i) {
+        threads.emplace_back([&] {
+            omp_set_num_threads(1);
+            int32_t distance;
+            faiss::idx_t label;
+            for (size_t j = 0; j < searches_per_thread; ++j) {
+                index.search(1, &database, 1, &distance, &label);
+            }
+        });
+    }
+    for (auto& thread : threads) {
+        thread.join();
+    }
+
+    EXPECT_EQ(num_threads * searches_per_thread, faiss::hnsw_stats.n1);
+}
 
 int reference_pop_min(faiss::MinimaxHeap& heap, float* vmin_out) {
     assert(heap.k > 0);