Add optional persistent locks to IndexHNSW for incremental adds

Summary:
`IndexHNSW` allocates an initializes locks for `ntotal+n` nodes on every call to `add()`. This makes batched insertion very costly, and incremental insertion prohibitively so.

This diff introduces optional persistent locks for `IndexHNSW` to improve incremental `add()` performance. Previously, `omp_lock_t` arrays of size `ntotal+n` were created/destroyed on each `add()` call. Now locks can be retained via a new `retain_locks` flag (default: false), using a new `HNSW::Lock` RAII wrapper with geometric growth.

RFC: Instead of `retain_locks` being the only way to opt into this new behavior, this could be inferred on the first incremental add. That is, clear the locks after insertion iff `n0 == 0`. Workloads which call `add()` once would be unaffected, but workloads which call `add()` repeatedly would 1) forego the clearing of the lock vector after `add()` call #2, and reuse locks for all subsequent calls. The downside would be the lack of the ability to reclaim the locks after insertion without HNSW-specific behavior at the call site.

Differential Revision: D98232750

Author: ddrcoder

benchs/bench_hnsw.py

@@ -189,3 +189,34 @@ def evaluate(index):
         print("search_L", search_L, end=' ')
         index.nsg.search_L = search_L
         evaluate(index)
+
+
+if "hnsw_locks" in todo:
+
+    batch_size = xb.shape[0] // 100
+    print(
+        f"Testing HNSW Flat: add with {batch_size=}, "
+        "with and without retaining locks"
+    )
+
+    for retain_locks in [False, True]:
+        index = faiss.IndexHNSWFlat(d, 32)
+        index.hnsw.efConstruction = 40
+        index.retain_locks = retain_locks
+
+        t0 = time.time()
+        t1 = None
+        t2 = None
+        for i in range(0, len(xb), batch_size):
+            t1 = time.time()
+            index.add(xb[i : i + batch_size])
+            t2 = time.time()
+            if i > 2 and t2 - t0 > 2:
+                break
+
+        assert t1 and t2
+        dt = t2 - t0
+        print(
+            f"\t {retain_locks=:1}: {index.ntotal} added in {t2 - t0:6.3f}s"
+            f" = {index.ntotal/(t2-t0):.0f}/s"
+        )

faiss/IndexBinaryHNSW.cpp

@@ -61,10 +61,7 @@ void hnsw_add_vertices(
         printf("  max_level = %d\n", max_level);
     }
 
-    std::vector<omp_lock_t> locks(ntotal);
-    for (size_t i = 0; i < ntotal; i++) {
-        omp_init_lock(&locks[i]);
-    }
+    HNSW::Locks locks(ntotal);
 
     // add vectors from highest to lowest level
     std::vector<int> hist;
@@ -156,10 +153,6 @@ void hnsw_add_vertices(
     if (verbose) {
         printf("Done in %.3f ms\n", getmillisecs() - t0);
     }
-
-    for (size_t i = 0; i < ntotal; i++) {
-        omp_destroy_lock(&locks[i]);
-    }
 }
 
 } // anonymous namespace

faiss/IndexHNSW.cpp

@@ -82,10 +82,8 @@ void hnsw_add_vertices(
         printf("  max_level = %d\n", max_level);
     }
 
-    std::vector<omp_lock_t> locks(ntotal);
-    for (size_t i = 0; i < ntotal; i++) {
-        omp_init_lock(&locks[i]);
-    }
+    auto& locks = index_hnsw.locks;
+    locks.resize(ntotal);
 
     // add vectors from highest to lowest level
     std::vector<int> hist;
@@ -199,9 +197,8 @@ void hnsw_add_vertices(
     if (verbose) {
         printf("Done in %.3f ms\n", getmillisecs() - t0);
     }
-
-    for (size_t i = 0; i < ntotal; i++) {
-        omp_destroy_lock(&locks[i]);
+    if (!index_hnsw.retain_locks) {
+        locks = HNSW::Locks{};
     }
 }
 
@@ -366,6 +363,7 @@ void IndexHNSW::reset() {
     hnsw.reset();
     storage->reset();
     ntotal = 0;
+    locks = HNSW::Locks{};
 }
 
 void IndexHNSW::reconstruct(idx_t key, float* recons) const {
@@ -526,10 +524,7 @@ void IndexHNSW::init_level_0_from_entry_points(
         int n,
         const storage_idx_t* points,
         const storage_idx_t* nearests) {
-    std::vector<omp_lock_t> locks(ntotal);
-    for (idx_t i = 0; i < ntotal; i++) {
-        omp_init_lock(&locks[i]);
-    }
+    locks.resize(ntotal);
 
 #pragma omp parallel
     {
@@ -547,7 +542,7 @@ void IndexHNSW::init_level_0_from_entry_points(
             dis->set_query(vec.data());
 
             hnsw.add_links_starting_from(
-                    *dis, pt_id, nearest, (*dis)(nearest), 0, locks.data(), vt);
+                    *dis, pt_id, nearest, (*dis)(nearest), 0, locks, vt);
 
             if (verbose && i % 10000 == 0) {
                 printf("  %d / %d\r", i, n);
@@ -559,8 +554,9 @@ void IndexHNSW::init_level_0_from_entry_points(
         printf("\n");
     }
 
-    for (idx_t i = 0; i < ntotal; i++) {
-        omp_destroy_lock(&locks[i]);
+    if (!retain_locks) {
+        locks = HNSW::Locks{};
+        FAISS_THROW_IF_NOT(locks.capacity() == 0);
     }
 }
 

faiss/IndexHNSW.h

@@ -52,11 +52,17 @@ struct IndexHNSW : Index {
     // See impl/VisitedTable.h.
     std::optional<bool> use_visited_hashset;
 
+    // Per-node locks for HNSW graph construction.
+    HNSW::Locks locks;
+    // locks are freed after each call to add() unless this flag is set.
+    bool retain_locks = false;
+
     explicit IndexHNSW(int d = 0, int M = 32, MetricType metric = METRIC_L2);
     explicit IndexHNSW(Index* storage, int M = 32);
 
     ~IndexHNSW() override;
 
+    /// Adds vectors to the index. May not be called concurrently.
     void add(idx_t n, const float* x) override;
 
     /// Trains the storage if needed

faiss/impl/HNSW.cpp

@@ -9,6 +9,7 @@
 
 #include <cinttypes>
 #include <cstddef>
+#include <cstdlib>
 
 #include <faiss/IndexHNSW.h>
 
@@ -504,7 +505,7 @@ void HNSW::add_links_starting_from(
         storage_idx_t nearest,
         float d_nearest,
         int level,
-        omp_lock_t* locks,
+        Locks& locks,
         VisitedTable& vt,
         bool keep_max_size_level0) {
     std::priority_queue<NodeDistCloser> link_targets;
@@ -543,7 +544,7 @@ void HNSW::add_with_locks(
         DistanceComputer& ptdis,
         int pt_level,
         int pt_id,
-        std::vector<omp_lock_t>& locks,
+        Locks& locks,
         VisitedTable& vt,
         bool keep_max_size_level0) {
     storage_idx_t nearest = entry_point;
@@ -580,7 +581,7 @@ void HNSW::add_with_locks(
                 nearest,
                 d_nearest,
                 level,
-                locks.data(),
+                locks,
                 vt,
                 keep_max_size_level0);
     }

faiss/impl/HNSW.h

@@ -17,7 +17,6 @@
 #include <faiss/impl/DistanceComputer.h>
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/maybe_owned_vector.h>
-#include <faiss/impl/platform_macros.h>
 #include <faiss/utils/Heap.h>
 #include <faiss/utils/random.h>
 
@@ -91,6 +90,26 @@ struct HNSW {
         int count_below(float thresh);
     };
 
+    /// RAII wrapper for `omp_lock_t`.
+    struct Lock : omp_lock_t {
+        Lock() {
+            omp_init_lock(this);
+        }
+        ~Lock() {
+            omp_destroy_lock(this);
+        }
+        // Copy is a no-op, defined to enable std::vector resize.
+        // Lock is assumed to not be held.
+        Lock(const Lock&) : Lock() {}
+        Lock& operator=(const Lock&) {
+            return *this;
+        }
+
+        Lock(Lock&&) = delete;
+        Lock& operator=(Lock&&) = delete;
+    };
+    using Locks = std::vector<Lock>;
+
     /// to sort pairs of (id, distance) from nearest to farthest or the reverse
     struct NodeDistCloser {
         float d;
@@ -196,7 +215,7 @@ struct HNSW {
             storage_idx_t nearest,
             float d_nearest,
             int level,
-            omp_lock_t* locks,
+            Locks& locks,
             VisitedTable& vt,
             bool keep_max_size_level0 = false);
 
@@ -206,7 +225,7 @@ struct HNSW {
             DistanceComputer& ptdis,
             int pt_level,
             int pt_id,
-            std::vector<omp_lock_t>& locks,
+            Locks& locks,
             VisitedTable& vt,
             bool keep_max_size_level0 = false);
 

faiss/python/swigfaiss.swig

@@ -557,6 +557,7 @@ void gpu_sync_all_devices()
 %include  <faiss/IndexIVFSpectralHash.h>
 %include  <faiss/IndexIVFAdditiveQuantizer.h>
 %include  <faiss/impl/HNSW.h>
+%ignore faiss::IndexHNSW::locks;
 %include  <faiss/IndexHNSW.h>
 
 %include <faiss/impl/kmeans1d.h>