Reuse the HNSW visited table across searches instead of reallocating per call (#5448)

Summary:
Pull Request resolved: #5448

## Background
Every `IndexHNSW::search()` call enters an `#pragma omp parallel` region in which each thread constructs its own `VisitedTable` via `VisitedTable::create(ntotal, ...)`. For the versioned-array strategy (`VisitedTableVector`), a fresh array allocation is paid **per thread, per Search() call** - this is an O(ntotal) allocation plus zero-fill overhead. It's freed every time the region exits.

## Problem
When a statically-built index is searched repeatedly with small or even single-query batches, this per-call alloc+zero can dominate the actual graph traversal (the search visits only a handful of nodes, but still pays to allocate and zero the whole `ntotal`-byte array).

## Solution
This adds `VisitedTable::get_reusable()` API to return a reference to the `thread_local` VisitedTable. The O(size) versioned array is **allocated once per thread and reused across all subsequent Search() calls**. This changed is applied to 3 code paths for search (`IndexHNSW.cpp`):
- `hnsw_search`
- `search_level_0` - only used by IndexHNSWCagra
- `IndexHNSWCagra::range_search` - only used by IndexHNSWCagra

Why `thread_local` rather than a member of the index:
- It is inherently per-OS-thread, so it survives across `search()` calls (OpenMP reuses its worker pool) and each thread gets its own table with no cross-thread coordination.
- faiss supports calling `const` search concurrently on one index from multiple threads. A member array indexed by `omp_get_thread_num()` would let two concurrent searches race on slot 0; `thread_local` cannot, since each OS thread has its own copy.
- The table grows on demand (`ensure_size`) if `ntotal` increases and never shrinks. Retained memory is bounded by the array-vs-hash-set threshold (D112025912).

## Correctness
`get_reusable()` calls `advance()` before returning the table. A prior search that threw exception without correctly `advance()` may have left version stamps NOT updated; without the reset those would read as spurious "visited" hits on the next search on that thread. `advance()` - might be redundant - in `get_reusable()` guarantees correct query version id.

## Minor notes
The `IndexHNSW2Level` mixed-search path is deliberately left on per-search `create()` rather than `get_reusable()`. It is a legacy path that uses the tri-state visited flags of `search_from_candidates_2` (two `advance()` calls per query), whose reset semantics differ from the bi-state paths above; keeping a fresh per-search table there leaves its behavior unchanged and scopes reuse to the paths where a single `advance()` at handout is correct.

Reviewed By: mnorris11, pankajsingh88

Differential Revision: D112042940

fbshipit-source-id: c6be0effdb2f52d21f6bb16db2d2ad0972a71bc0

Author: Lei Huang

faiss/IndexHNSW.cpp

@@ -280,12 +280,12 @@ void hnsw_search(
 
 #pragma omp parallel if (i1 - i0 > 1)
         {
-            std::unique_ptr<VisitedTable> vt;
+            VisitedTable* vt = nullptr;
             std::unique_ptr<typename BlockResultHandler::SingleResultHandler>
                     res;
             std::unique_ptr<DistanceComputer> dis;
             try {
-                vt = VisitedTable::create(
+                vt = &VisitedTable::get_reusable(
                         index->ntotal, hnsw.use_visited_hashset);
                 res = std::make_unique<
                         typename BlockResultHandler::SingleResultHandler>(bres);
@@ -479,11 +479,11 @@ void IndexHNSW::search_level_0(
         {
             std::unique_ptr<DistanceComputer> qdis;
             HNSWStats search_stats;
-            std::unique_ptr<VisitedTable> vt;
+            VisitedTable* vt = nullptr;
             std::unique_ptr<typename RH::SingleResultHandler> res;
             try {
                 qdis.reset(storage_distance_computer(storage));
-                vt = VisitedTable::create(
+                vt = &VisitedTable::get_reusable(
                         hnsw_ntotal, hnsw.use_visited_hashset);
                 res = std::make_unique<typename RH::SingleResultHandler>(bres);
             } catch (...) {
@@ -1208,11 +1208,11 @@ void IndexHNSWCagra::range_search(
 
             RangeQueryResult& qres = pres.new_result(i);
             RangeResultHandler<C> res(&qres, radius);
-            std::unique_ptr<VisitedTable> vt =
-                    VisitedTable::create(ntotal, hnsw.use_visited_hashset);
+            VisitedTable& vt = VisitedTable::get_reusable(
+                    ntotal, hnsw.use_visited_hashset);
             HNSWStats stats;
             hnsw.search_level_0(
-                    *dis, res, 1, &nearest, &nearest_d, 1, stats, *vt, params);
+                    *dis, res, 1, &nearest, &nearest_d, 1, stats, vt, params);
             n1 += stats.n1;
             n2 += stats.n2;
             ndis += stats.ndis;

faiss/impl/VisitedTable.cpp

@@ -33,6 +33,22 @@ std::unique_ptr<VisitedTable> VisitedTable::create(
     return std::make_unique<VisitedTableVector>(size);
 }
 
+VisitedTable& VisitedTable::get_reusable(
+        size_t size,
+        std::optional<bool> use_hashset) {
+    bool use_set =
+            use_hashset.value_or(size >= visited_table_hashset_threshold);
+    if (use_set) {
+        thread_local VisitedTableSet tls_set;
+        tls_set.advance();
+        return tls_set;
+    }
+    thread_local VisitedTableVector tls_vec(0);
+    tls_vec.ensure_size(size);
+    tls_vec.advance();
+    return tls_vec;
+}
+
 void VisitedTableVector::advance() {
     if (visno < 254) {
         // 254 rather than 255 because sometimes we use visno and visno+1

faiss/impl/VisitedTable.h

@@ -48,6 +48,18 @@ struct VisitedTable {
     static std::unique_ptr<VisitedTable> create(
             size_t size,
             std::optional<bool> use_hashset = std::nullopt);
+
+    /// Returns a thread-local, reusable table sized for at least `size` and
+    /// reset to a clean state. Unlike create(), it does not allocate on each
+    /// call: the O(size) versioned array is allocated once per thread and
+    /// reused across searches, avoiding a per-search alloc+zero of the whole
+    /// array when a static index is searched repeatedly.
+    ///
+    /// The returned reference is owned by thread-local storage: do not delete
+    /// it and do not use it beyond the current search on the calling thread.
+    static VisitedTable& get_reusable(
+            size_t size,
+            std::optional<bool> use_hashset = std::nullopt);
 };
 
 /// Set-based implementation using unordered_set.
@@ -87,6 +99,14 @@ struct VisitedTableVector FAISS_FINAL : VisitedTable {
 
     explicit VisitedTableVector(size_t size) : visited(size, 0) {}
 
+    /// Grow so indices in [0, size) are valid; new slots read as unvisited.
+    /// Never shrinks, so capacity is retained when the table is reused.
+    void ensure_size(size_t size) {
+        if (visited.size() < size) {
+            visited.resize(size, 0);
+        }
+    }
+
     bool set(size_t no) final {
         if (visited[no] == visno) {
             return false;

tests/test_hnsw.cpp

@@ -750,3 +750,73 @@ TEST_F(HNSWTest, TEST_search_level_0) {
     EXPECT_GT(stats1.n1, stats2.n1);
     EXPECT_GT(stats1.n2, stats2.n2);
 }
+
+TEST(VisitedTableReuse, ReusesGrowsAndResets) {
+    // get_reusable() hands back the same thread-local object across calls
+    // (no per-call allocation) ...
+    faiss::VisitedTable& t1 = faiss::VisitedTable::get_reusable(100);
+    faiss::VisitedTable& t2 = faiss::VisitedTable::get_reusable(100);
+    EXPECT_EQ(&t1, &t2);
+
+    // ... and hands it back clean: a mark set before re-acquiring must not
+    // survive, since a reused table may carry stamps from a prior search.
+    t2.set(42);
+    EXPECT_TRUE(t2.get(42));
+    faiss::VisitedTable& t3 = faiss::VisitedTable::get_reusable(100);
+    EXPECT_EQ(&t2, &t3);
+    EXPECT_FALSE(t3.get(42));
+
+    // Growing to a larger size keeps the same object and exposes the new range.
+    faiss::VisitedTable& t4 = faiss::VisitedTable::get_reusable(1000);
+    EXPECT_EQ(&t3, &t4);
+    EXPECT_FALSE(t4.get(900));
+    EXPECT_TRUE(t4.set(900));
+    EXPECT_TRUE(t4.get(900));
+}
+
+// Restores the process-global OpenMP thread count on scope exit, so a test that
+// changes it via omp_set_num_threads does not leak the setting into later
+// tests.
+struct ScopedOmpThreads {
+    int saved = omp_get_max_threads();
+    ~ScopedOmpThreads() {
+        omp_set_num_threads(saved);
+    }
+};
+
+TEST_F(HNSWTest, TEST_search_reuse_correctness) {
+    ScopedOmpThreads omp_guard;
+
+    std::vector<faiss::idx_t> I1(k * nq), I2(k * nq);
+    std::vector<float> D1(k * nq), D2(k * nq);
+
+    // Two back-to-back searches must be identical: the reused visited table
+    // must not carry state across search() calls.
+    omp_set_num_threads(1);
+    index->search(nq, xq->data(), k, D1.data(), I1.data());
+    index->search(nq, xq->data(), k, D2.data(), I2.data());
+    EXPECT_EQ(I1, I2);
+    EXPECT_EQ(D1, D2);
+
+    // The reused versioned-array path must match a hash-set reference.
+    std::vector<faiss::idx_t> Iref(k * nq);
+    std::vector<float> Dref(k * nq);
+    index->hnsw.use_visited_hashset = true;
+    index->search(nq, xq->data(), k, Dref.data(), Iref.data());
+    index->hnsw.use_visited_hashset = false;
+    std::vector<faiss::idx_t> Iarr(k * nq);
+    std::vector<float> Darr(k * nq);
+    index->search(nq, xq->data(), k, Darr.data(), Iarr.data());
+    EXPECT_EQ(Iref, Iarr);
+    EXPECT_EQ(Dref, Darr);
+
+    // Multi-threaded search must match single-threaded: each thread reuses its
+    // own thread-local table.
+    index->hnsw.use_visited_hashset = std::nullopt;
+    omp_set_num_threads(4);
+    std::vector<faiss::idx_t> Imt(k * nq);
+    std::vector<float> Dmt(k * nq);
+    index->search(nq, xq->data(), k, Dmt.data(), Imt.data());
+    EXPECT_EQ(I1, Imt);
+    EXPECT_EQ(D1, Dmt);
+}