Make search stats collection configurable (#5670)

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

Follow-up to D121368933 and public PR https://github.com/facebookresearch/faiss/pull/5657.

HNSW and RaBitQ search statistics are available to C++ and Python. Faiss collects them by default. Every search updates shared counters. Under high concurrency, these updates are expensive. In the 64-thread relaxed-atomic profile, `HNSWStats::combine_atomic` uses 74.14% of the sampled user cycles.

This diff adds one process-wide flag in `faiss/utils/utils.h`:

- `set_search_stats_enabled(bool)`
- `get_search_stats_enabled()`

The flag is a relaxed `std::atomic<bool>`. It is private to `utils.cpp`. The default is enabled, for backward compatibility. Callers that do not read the statistics can disable collection. Then the searches skip the global atomic updates.

At the moment, the flag controls `hnsw_stats` and `rabitq_stats`. The name is generic on purpose. The other global stats objects (for example `indexIVF_stats`, `indexPQ_stats` and `indexPanorama_stats`) use plain `+=` on shared counters, so they have the same scaling problem. Later diffs can put these objects under the same flag without a new API. The documentation says that the values of global stats objects are unspecified when the flag is disabled. Thus, more coverage does not change the contract.

API design:

- Faiss often exposes a plain global for a setting, for example `distance_compute_blas_threshold`. This diff does not do that. The flag is for concurrent searches, and a thread can change it while other threads search. A plain `bool` is then a data race, and TSAN reports it. This stack fixed TSAN races in D121097648.
- SWIG cannot wrap `std::atomic<bool>`. Thus, the setter and the getter are the only API. Python calls `faiss.set_search_stats_enabled(False)`.

Other changes:

- `bench_hnsw_stats` has a `stats` argument. `stats:0` disables collection, and `stats:1` enables it. One binary measures the two modes.
- The tests save the flag value and restore it after the test. The C++ tests use `SearchStatsEnabledGuard`. The Python tests use `addCleanup`. Thus, a failed assertion does not change the flag for later tests in the same process.

Mode/opt benchmark from V12, five sequential repetitions (mean real-time throughput):

| External pthreads | Relaxed atomic stats enabled | Stats disabled | Change |
| ---: | ---: | ---: | ---: |
| 1 | 2.516 M/s | 2.411 M/s | 0.96x |
| 2 | 3.919 M/s | 4.164 M/s | 1.06x |
| 4 | 7.494 M/s | 8.260 M/s | 1.10x |
| 8 | 12.717 M/s | 14.170 M/s | 1.11x |
| 16 | 16.074 M/s | 22.979 M/s | 1.43x |
| 32 | 19.929 M/s | 36.841 M/s | 1.85x |
| 64 | 22.101 M/s | 66.642 M/s | 3.02x |

The 0.96x value at 1 pthread is noise. A longer pinned measurement in the test plan shows no slowdown. At 1 thread, the atomic adds have no contention, so the expected gain is small.

The disabled path scales 27.64x from 1 to 64 pthreads. The enabled path scales 8.78x. The disabled path does not scale linearly, because this one-vector search is very small. At 64 threads, the OpenMP runtime uses about 28% of the sampled cycles. Allocator, RTTI and shared-cache costs also contribute.

Reviewed By: luciang

Differential Revision: D121374405

fbshipit-source-id: a54ee81a68a63df01c93bd5098171658883cd5ab

benchs/bench_hnsw_stats.cpp

@@ -11,6 +11,7 @@
 
 #include <benchmark/benchmark.h>
 #include <faiss/IndexBinaryHNSW.h>
+#include <faiss/utils/utils.h>
 #include <omp.h>
 
 namespace faiss {
@@ -41,6 +42,7 @@ const IndexBinaryHNSW& get_index() {
 
 void bench_binary_hnsw_single_query(benchmark::State& state) {
     omp_set_num_threads(1);
+    set_search_stats_enabled(state.range(0) != 0);
     constexpr std::array<uint8_t, kDimension / 8> query{};
     int32_t distance;
     idx_t label;
@@ -53,7 +55,12 @@ void bench_binary_hnsw_single_query(benchmark::State& state) {
     state.SetItemsProcessed(state.iterations());
 }
 
-BENCHMARK(bench_binary_hnsw_single_query)->ThreadRange(1, 64)->UseRealTime();
+BENCHMARK(bench_binary_hnsw_single_query)
+        ->ArgName("stats")
+        ->Arg(0)
+        ->Arg(1)
+        ->ThreadRange(1, 64)
+        ->UseRealTime();
 
 } // namespace
 } // namespace faiss

faiss/IndexBinaryHNSW.cpp

@@ -24,6 +24,7 @@
 #include <faiss/impl/VisitedTable.h>
 #include <faiss/utils/Heap.h>
 #include <faiss/utils/random.h>
+#include <faiss/utils/utils.h>
 
 #include <faiss/impl/simd_dispatch.h>
 
@@ -121,7 +122,9 @@ void IndexBinaryHNSW::search(
         }
     }
 
-    hnsw_stats.combine_atomic({n1, n2, ndis, nhops});
+    if (get_search_stats_enabled()) {
+        hnsw_stats.combine_atomic({n1, n2, ndis, nhops});
+    }
 
 #pragma omp parallel for
     for (idx_t i = 0; i < n * k; ++i) {
@@ -270,7 +273,9 @@ void IndexBinaryHNSWCagra::search(
 
                 res.end();
             }
-            hnsw_stats.combine_atomic(search_stats);
+            if (get_search_stats_enabled()) {
+                hnsw_stats.combine_atomic(search_stats);
+            }
         }
 
 #pragma omp parallel for

faiss/IndexHNSW.cpp

@@ -35,6 +35,7 @@
 #include <faiss/impl/platform_macros.h>
 #include <faiss/utils/random.h>
 #include <faiss/utils/sorting.h>
+#include <faiss/utils/utils.h>
 
 namespace faiss {
 
@@ -495,8 +496,10 @@ void hnsw_search(
         InterruptCallback::check();
     }
 
-    hnsw_stats.combine_atomic({n1, n2, ndis, nhops});
-    rabitq_stats.add_atomic({n_rabitq_1bit, n_rabitq_refine});
+    if (get_search_stats_enabled()) {
+        hnsw_stats.combine_atomic({n1, n2, ndis, nhops});
+        rabitq_stats.add_atomic({n_rabitq_1bit, n_rabitq_refine});
+    }
 }
 
 } // anonymous namespace
@@ -707,8 +710,10 @@ void IndexHNSW::search_level_0(
                     omp_capture_exception(ex, [&] { interrupt = true; });
                 }
             }
-            hnsw_stats.combine_atomic(search_stats);
-            rabitq_stats.add_atomic(rq_search_stats);
+            if (get_search_stats_enabled()) {
+                hnsw_stats.combine_atomic(search_stats);
+                rabitq_stats.add_atomic(rq_search_stats);
+            }
         }
         omp_rethrow_if_exception(ex);
     };
@@ -1235,7 +1240,9 @@ void IndexHNSW2Level::search(
         }
         omp_rethrow_if_exception(ex);
 
-        hnsw_stats.combine_atomic({n1, n2, ndis, nhops});
+        if (get_search_stats_enabled()) {
+            hnsw_stats.combine_atomic({n1, n2, ndis, nhops});
+        }
     }
 }
 
@@ -1437,7 +1444,9 @@ void IndexHNSWCagra::range_search(
         result->do_allocation();
         pres.copy_result();
 
-        hnsw_stats.combine_atomic({n1, n2, ndis, nhops});
+        if (get_search_stats_enabled()) {
+            hnsw_stats.combine_atomic({n1, n2, ndis, nhops});
+        }
     };
 
     if (is_similarity_metric(metric_type)) {

faiss/python/__init__.pyi

@@ -3114,6 +3114,8 @@ class SIMDConfig:
 # Utility functions
 def get_mem_usage_kb() -> int: ...
 def get_compile_options() -> str: ...
+def set_search_stats_enabled(enabled: bool) -> None: ...
+def get_search_stats_enabled() -> bool: ...
 def check_openmp() -> bool: ...
 def shard_ivf_index_centroids(
     index: IndexIVF,

faiss/utils/utils.cpp

@@ -28,6 +28,7 @@
 #include <omp.h>
 
 #include <algorithm>
+#include <atomic>
 #include <set>
 #include <type_traits>
 #include <unordered_set>
@@ -175,6 +176,18 @@ uint64_t get_cycles() {
 #endif
 }
 
+namespace {
+std::atomic<bool> search_stats_enabled{true};
+} // namespace
+
+void set_search_stats_enabled(bool enabled) {
+    search_stats_enabled.store(enabled, std::memory_order_relaxed);
+}
+
+bool get_search_stats_enabled() {
+    return search_stats_enabled.load(std::memory_order_relaxed);
+}
+
 #ifdef __linux__
 
 size_t get_mem_usage_kb() {

faiss/utils/utils.h

@@ -48,6 +48,14 @@ size_t get_mem_usage_kb();
 
 uint64_t get_cycles();
 
+/** Enable or disable the collection of global search statistics (enabled by
+ * default). The flag is process-wide. When it is disabled, searches may skip
+ * updates to global stats objects, so their values are unspecified. At the
+ * moment, the flag controls hnsw_stats and rabitq_stats.
+ */
+void set_search_stats_enabled(bool enabled);
+bool get_search_stats_enabled();
+
 /***************************************************************************
  * Misc  matrix and vector manipulation functions
  ***************************************************************************/

tests/test_graph_based.py

@@ -217,6 +217,10 @@ class TestHNSW(unittest.TestCase):
         index = faiss.IndexHNSWFlat(d, 16)
         index.add(self.xb)
         stats = faiss.cvar.hnsw_stats
+        self.addCleanup(
+            faiss.set_search_stats_enabled, faiss.get_search_stats_enabled()
+        )
+        faiss.set_search_stats_enabled(True)
         stats.reset()
         Dhnsw, Ihnsw = index.search(self.xq, 1)
         self.assertGreater(stats.ndis, len(self.xq) * index.hnsw.efSearch)

tests/test_hnsw.cpp

@@ -24,9 +24,40 @@
 #include <faiss/impl/VisitedTable.h>
 #include <faiss/impl/hnsw/MinimaxHeap.h>
 #include <faiss/utils/random.h>
+#include <faiss/utils/utils.h>
 #include <omp.h>
 
-TEST(HNSWStats, ConcurrentBinarySearchAggregatesExactly) {
+namespace {
+
+struct SearchStatsEnabledGuard {
+    explicit SearchStatsEnabledGuard(bool enabled)
+            : previous(faiss::get_search_stats_enabled()) {
+        faiss::set_search_stats_enabled(enabled);
+    }
+    ~SearchStatsEnabledGuard() {
+        faiss::set_search_stats_enabled(previous);
+    }
+    bool previous;
+};
+
+} // namespace
+
+TEST(HNSWStats, DisabledCollectionSkipsGlobalStats) {
+    faiss::IndexBinaryHNSW index(8, 2);
+    const uint8_t database = 0;
+    index.add(1, &database);
+    index.hnsw.efSearch = 1;
+    SearchStatsEnabledGuard guard(false);
+    faiss::hnsw_stats.reset();
+
+    int32_t distance;
+    faiss::idx_t label;
+    index.search(1, &database, 1, &distance, &label);
+
+    EXPECT_EQ(0, faiss::hnsw_stats.n1);
+}
+
+TEST(HNSWStats, EnabledCollectionAggregatesConcurrentSearchesExactly) {
     constexpr size_t num_threads = 16;
     constexpr size_t searches_per_thread = 100;
 
@@ -34,6 +65,7 @@ TEST(HNSWStats, ConcurrentBinarySearchAggregatesExactly) {
     const uint8_t database = 0;
     index.add(1, &database);
     index.hnsw.efSearch = 1;
+    SearchStatsEnabledGuard guard(true);
     faiss::hnsw_stats.reset();
 
     std::vector<std::thread> threads;

tests/test_hnsw_rabitq.py

@@ -10,6 +10,11 @@ import numpy as np
 
 
 class TestHNSWRaBitQ(unittest.TestCase):
+    def setUp(self):
+        self.addCleanup(
+            faiss.set_search_stats_enabled, faiss.get_search_stats_enabled()
+        )
+
     def make_data(self, d=32, nt=400, nb=600, nq=20):
         rs = np.random.RandomState(123)
         xt = rs.randn(nt, d).astype("float32")
@@ -39,6 +44,13 @@ class TestHNSWRaBitQ(unittest.TestCase):
     def test_staged_search_quality(self):
         index, xb, xq = self.make_index(nb_bits=3)
 
+        faiss.set_search_stats_enabled(False)
+        faiss.cvar.rabitq_stats.reset()
+        index.search(xq, 10)
+        self.assertEqual(faiss.cvar.rabitq_stats.n_1bit, 0)
+        self.assertEqual(faiss.cvar.rabitq_stats.n_refine, 0)
+
+        faiss.set_search_stats_enabled(True)
         faiss.cvar.rabitq_stats.reset()
         D, I = index.search(xq, 10)
         stats = faiss.cvar.rabitq_stats
@@ -130,6 +142,7 @@ class TestHNSWRaBitQ(unittest.TestCase):
     def test_one_bit_fallback(self):
         index, _, xq = self.make_index(nb_bits=1)
         self.assertEqual(index.hnsw.search_method, 0)
+        faiss.set_search_stats_enabled(True)
         faiss.cvar.rabitq_stats.reset()
         D, I = index.search(xq, 10)
         self.assertTrue(np.all(I >= 0))