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

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))