diff --git a/benchs/bench_hnsw_stats.cpp b/benchs/bench_hnsw_stats.cpp index 4aba3feb8..32968d59c 100644 --- a/benchs/bench_hnsw_stats.cpp +++ b/benchs/bench_hnsw_stats.cpp @@ -11,6 +11,7 @@ #include #include +#include #include 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 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 diff --git a/faiss/IndexBinaryHNSW.cpp b/faiss/IndexBinaryHNSW.cpp index f3b9f11fb..3dd738bde 100644 --- a/faiss/IndexBinaryHNSW.cpp +++ b/faiss/IndexBinaryHNSW.cpp @@ -24,6 +24,7 @@ #include #include #include +#include #include @@ -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 diff --git a/faiss/IndexHNSW.cpp b/faiss/IndexHNSW.cpp index cf4885c0f..d295a3dd4 100644 --- a/faiss/IndexHNSW.cpp +++ b/faiss/IndexHNSW.cpp @@ -35,6 +35,7 @@ #include #include #include +#include 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)) { diff --git a/faiss/python/__init__.pyi b/faiss/python/__init__.pyi index 152349526..8ebb208ed 100644 --- a/faiss/python/__init__.pyi +++ b/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, diff --git a/faiss/utils/utils.cpp b/faiss/utils/utils.cpp index 7a14fc2cf..abd4fc756 100644 --- a/faiss/utils/utils.cpp +++ b/faiss/utils/utils.cpp @@ -28,6 +28,7 @@ #include #include +#include #include #include #include @@ -175,6 +176,18 @@ uint64_t get_cycles() { #endif } +namespace { +std::atomic 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() { diff --git a/faiss/utils/utils.h b/faiss/utils/utils.h index 3af2590dc..e387c6263 100644 --- a/faiss/utils/utils.h +++ b/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 ***************************************************************************/ diff --git a/tests/test_graph_based.py b/tests/test_graph_based.py index f314aa41a..3a659897e 100644 --- a/tests/test_graph_based.py +++ b/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) diff --git a/tests/test_hnsw.cpp b/tests/test_hnsw.cpp index a0be2193e..163ccadce 100644 --- a/tests/test_hnsw.cpp +++ b/tests/test_hnsw.cpp @@ -24,9 +24,40 @@ #include #include #include +#include #include -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 threads; diff --git a/tests/test_hnsw_rabitq.py b/tests/test_hnsw_rabitq.py index f338a1863..4e3ce9a2c 100644 --- a/tests/test_hnsw_rabitq.py +++ b/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))