facebookresearch
diff --git a/‎faiss/IndexHNSW.cpp‎
Lines changed: 239 additions & 67 deletions b/‎faiss/IndexHNSW.cpp‎
Lines changed: 239 additions & 67 deletions
diff --git a/‎faiss/IndexHNSW.h‎
Lines changed: 2 additions & 1 deletion b/‎faiss/IndexHNSW.h‎
Lines changed: 2 additions & 1 deletion
@@ -60,7 +60,14 @@ DistanceComputer* storage_distance_computer(const Index* storage) {
     return storage->get_distance_computer();
 }
 
-void hnsw_add_vertices(
+// The HNSW graph build: deterministic and lock-free. Points are bucketed by
+// level and shuffled, then each bucket is inserted in prefix-doubling batches.
+// Per batch, phase A computes forward links against the frozen snapshot and
+// records reciprocal edges; phase B groups those by destination and merges them
+// one node at a time in a fixed order. Reproducible across runs at a fixed
+// thread count. Supports the CAGRA level-0 import path (init_level0=false skips
+// level 0; keep_max_size_level0 fills the base layer to 2*M).
+void hnsw_add_vertices_deterministic(
         IndexHNSW& index_hnsw,
         size_t n0,
         size_t n,
@@ -72,7 +79,7 @@ void hnsw_add_vertices(
     size_t ntotal = n0 + n;
     double t0 = getmillisecs();
     if (verbose) {
-        printf("hnsw_add_vertices: adding %zd elements on top of %zd "
+        printf("hnsw_add_vertices_deterministic: adding %zd elements on top of %zd "
                "(preset_levels=%d)\n",
                n,
                n0,
@@ -83,16 +90,20 @@ void hnsw_add_vertices(
         return;
     }
 
-    int max_level = hnsw.prepare_level_tab(n, preset_levels);
+    // Lowest level *bucket* to insert. With init_level0=false (the CAGRA
+    // level-0 import path) the level-0-only points are skipped, exactly as the
+    // default build's `pt_level >= !init_level0` bound; points that live on
+    // higher levels are still inserted and still build their own level-0 links.
+    const int min_bucket_level = index_hnsw.init_level0 ? 0 : 1;
+    const bool keep_max_size_level0 = index_hnsw.keep_max_size_level0;
 
+    int max_level = hnsw.prepare_level_tab(n, preset_levels);
     if (verbose) {
         printf("  max_level = %d\n", max_level);
     }
 
-    auto& locks = index_hnsw.locks;
-    locks.prepare(ntotal);
-
-    // add vectors from highest to lowest level
+    // Bucket the new points by level (highest first), exactly as
+    // hnsw_add_vertices does.
     std::vector<int> hist;
     std::vector<int> order(n);
 
@@ -122,94 +133,260 @@ void hnsw_add_vertices(
         }
     }
 
-    idx_t check_period = InterruptCallback::get_period_hint(
+    // Upper bound on batch size (ParlayANN's theta), 2% of the index.
+    const size_t theta =
+            std::max<size_t>(1, static_cast<size_t>(0.02 * ntotal));
+
+    // Poll the interrupt callback roughly every check_period points inside the
+    // parallel forward-link phase (phase A), mirroring the default build. A
+    // single prefix-doubling batch can be up to theta = 2% of ntotal points, so
+    // without this the build would only honor cancellation at the batch
+    // boundary -- on very large indexes that can be minutes.
+    const idx_t check_period = InterruptCallback::get_period_hint(
             max_level * index_hnsw.d * hnsw.efConstruction);
 
-    { // perform add
-        RandomGenerator rng2(789);
+    RandomGenerator rng2(789);
+    size_t i1 = n;
+
+    for (int pt_level = static_cast<int>(hist.size()) - 1;
+         pt_level >= min_bucket_level;
+         pt_level--) {
+        size_t i0 = i1 - hist[pt_level];
+        if (i0 == i1) {
+            continue;
+        }
 
-        size_t i1 = static_cast<int>(n);
+        if (verbose) {
+            printf("Adding %zu elements at level %d\n", i1 - i0, pt_level);
+        }
 
-        for (int pt_level = static_cast<int>(hist.size()) - 1;
-             pt_level >= int(!index_hnsw.init_level0);
-             pt_level--) {
-            size_t i0 = i1 - hist[pt_level];
+        // random permutation to get rid of dataset order bias
+        for (size_t j = i0; j < i1; j++) {
+            std::swap(
+                    order[j],
+                    order[j + rng2.rand_int(static_cast<int>(i1 - j))]);
+        }
 
-            if (verbose) {
-                printf("Adding %zu elements at level %d\n", i1 - i0, pt_level);
-            }
+        // Deterministically bootstrap / raise the entry point. The top
+        // bucket is processed first; its first point (in shuffled order) is
+        // a valid maximum-level entry point.
+        if (hnsw.entry_point == -1 || pt_level > hnsw.max_level) {
+            hnsw.max_level = pt_level;
+            hnsw.entry_point = order[i0];
+        }
 
-            // random permutation to get rid of dataset order bias
-            for (size_t j = i0; j < i1; j++) {
-                std::swap(
-                        order[j],
-                        order[j + rng2.rand_int(static_cast<int>(i1 - j))]);
+        // Prefix-doubling batches within this bucket.
+        size_t s = i0;
+        while (s < i1) {
+            size_t done = s - i0;
+            size_t grow = std::min(done == 0 ? size_t(1) : done, theta);
+            size_t e = std::min(i1, s + grow);
+
+            // ---- Phase A: compute forward links against the snapshot ----
+            // Reciprocal edge recorded during phase A: destination `dest` gains
+            // an incoming link from `src` at level `level`.
+            struct Edge {
+                int level;
+                HNSW::storage_idx_t dest;
+                HNSW::storage_idx_t src;
+            };
+
+            // Reverse edges are bump-allocated directly into one pre-sized
+            // buffer during phase A -- there is no per-thread accumulator. The
+            // buffer is sized to the exact upper bound (the sum of the batch
+            // points' neighbor-slot counts); each point writes at most that
+            // many forward links, so the actual fill is <= cap_sum and the
+            // buffer never reallocates. `edge_counter` hands out packed
+            // contiguous runs; its (nondeterministic) ordering does not affect
+            // the graph because phase B re-groups by destination and re-sorts
+            // each node's incoming set.
+            size_t cap_sum = 0;
+            for (int64_t i = s; i < static_cast<int64_t>(e); i++) {
+                storage_idx_t pid = order[i];
+                cap_sum += hnsw.offsets[pid + 1] - hnsw.offsets[pid];
             }
-
+            std::vector<Edge> edges(cap_sum);
+            std::atomic<size_t> edge_counter{0};
+
+            // Set from inside the parallel region on interrupt and thrown after
+            // it (we cannot throw out of an OpenMP region). The unsynchronized
+            // write is a benign race, as in the default build: at worst the
+            // interrupt is noticed one poll late. It never changes the graph --
+            // an uninterrupted build processes every point identically.
             bool interrupt = false;
 
-#pragma omp parallel if (i1 > i0 + 100)
+#pragma omp parallel if (e - s > 100)
             {
                 std::unique_ptr<VisitedTable> vt =
                         VisitedTable::create(ntotal, hnsw.use_visited_hashset);
-
                 std::unique_ptr<DistanceComputer> dis(
                         storage_distance_computer(index_hnsw.storage));
-                bool do_display = verbose && omp_get_thread_num() == 0;
-                size_t prev_display = 0;
+
+                std::vector<std::pair<HNSW::storage_idx_t, int>> reverse_edges;
+
+                // Per-thread poll counter. With schedule(static) each thread
+                // walks a contiguous chunk, so counting our own iterations
+                // polls the interrupt callback every check_period points
+                // regardless of how the chunk aligns with the global index -- a
+                // global
+                // `(i - s) % check_period` would leave threads whose chunk
+                // straddles no multiple of check_period never polling.
                 size_t counter = 0;
 
-                // here we should do schedule(dynamic) but this segfaults for
-                // some versions of LLVM. The performance impact should not be
-                // too large when (i1 - i0) / num_threads >> 1
 #pragma omp for schedule(static)
-                for (int64_t i = i0; i < i1; i++) {
-                    storage_idx_t pt_id = order[i];
-                    dis->set_query(x + (pt_id - n0) * d);
-
-                    // cannot break
+                for (int64_t i = s; i < static_cast<int64_t>(e); i++) {
                     if (interrupt) {
-                        continue;
+                        continue; // cannot break out of an OpenMP for loop
                     }
+                    storage_idx_t pt_id = order[i];
+                    int lvl = hnsw.levels[pt_id] - 1;
+                    dis->set_query(x + (pt_id - n0) * d);
 
-                    hnsw.add_with_locks(
+                    reverse_edges.clear();
+                    hnsw.compute_forward_links_deterministic(
                             *dis,
-                            pt_level,
+                            lvl,
                             pt_id,
-                            locks,
                             *vt,
-                            index_hnsw.keep_max_size_level0 && (pt_level == 0));
-
-                    if (do_display && i - i0 > prev_display + 10000) {
-                        prev_display = i - i0;
-                        printf("  %zu / %zu\r", i - i0, i1 - i0);
-                        fflush(stdout);
+                            reverse_edges,
+                            keep_max_size_level0);
+
+                    // Grab a contiguous run in the buffer with one atomic bump,
+                    // then fill it (packed, no gaps).
+                    size_t off = edge_counter.fetch_add(
+                            reverse_edges.size(), std::memory_order_relaxed);
+                    for (size_t k = 0; k < reverse_edges.size(); k++) {
+                        edges[off + k] = {
+                                reverse_edges[k].second,
+                                reverse_edges[k].first,
+                                pt_id};
                     }
-                    if (counter % check_period == 0) {
-                        if (InterruptCallback::is_interrupted()) {
-                            interrupt = true;
-                        }
+
+                    if (counter++ % check_period == 0 &&
+                        InterruptCallback::is_interrupted()) {
+                        interrupt = true;
                     }
-                    counter++;
                 }
             }
             if (interrupt) {
                 FAISS_THROW_MSG("computation interrupted");
             }
-            i1 = i0;
-        }
-        if (index_hnsw.init_level0) {
-            FAISS_ASSERT(i1 == 0);
-        } else {
-            FAISS_ASSERT((i1 - hist[0]) == 0);
+
+            // ---- Phase B: deterministic reciprocal-edge merge ----
+            // Group the edges by (level, dest) via a fixed-size radix partition
+            // on the low bits of dest, then merge each bucket in parallel. The
+            // bucket count is a small constant (independent of ntotal and
+            // thread count), so this stays O(edges) in memory. Every dest maps
+            // to exactly one bucket, so each (level, dest) group is wholly
+            // contained in one bucket; the merge is therefore order- and
+            // thread-count-independent (merge_reverse_links_deterministic
+            // re-sorts and dedups each node's incoming set internally), which
+            // keeps the resulting graph byte-identical across builds. The radix
+            // partition runs in place, so `edges` stays the only edge buffer.
+            const size_t total = edge_counter.load();
+
+            constexpr int kBucketBits = 8;
+            constexpr uint32_t kNumBuckets = 1u << kBucketBits;
+            constexpr uint32_t kBucketMask = kNumBuckets - 1;
+
+            // bstart[b]..bstart[b+1] delimits bucket b after partitioning.
+            std::vector<size_t> bstart(kNumBuckets + 1, 0);
+            for (size_t idx = 0; idx < total; idx++) {
+                bstart[(static_cast<uint32_t>(edges[idx].dest) & kBucketMask) +
+                       1]++;
+            }
+            for (uint32_t b = 0; b < kNumBuckets; b++) {
+                bstart[b + 1] += bstart[b];
+            }
+
+            // In-place radix partition (cycle permutation): each iteration
+            // places at least one edge in its final bucket, so this is O(total)
+            // and needs no auxiliary edge buffer.
+            {
+                std::vector<size_t> head(bstart.begin(), bstart.end() - 1);
+                for (uint32_t b = 0; b < kNumBuckets; b++) {
+                    size_t end_b = bstart[b + 1];
+                    while (head[b] < end_b) {
+                        Edge cur_e = edges[head[b]];
+                        if ((static_cast<uint32_t>(cur_e.dest) & kBucketMask) ==
+                            b) {
+                            head[b]++;
+                            continue;
+                        }
+                        while ((static_cast<uint32_t>(cur_e.dest) &
+                                kBucketMask) != b) {
+                            uint32_t tb = static_cast<uint32_t>(cur_e.dest) &
+                                    kBucketMask;
+                            std::swap(cur_e, edges[head[tb]]);
+                            head[tb]++;
+                        }
+                        edges[head[b]] = cur_e;
+                        head[b]++;
+                    }
+                }
+            }
+            Edge* base = edges.data();
+
+#pragma omp parallel if (total > 100)
+            {
+                std::unique_ptr<DistanceComputer> dis(
+                        storage_distance_computer(index_hnsw.storage));
+                // schedule(static), not (dynamic): the LLVM/libomp dynamic
+                // dispatcher segfaults in some build configs (see the note in
+                // the pre-existing lock-based build). Buckets are ~uniform in
+                // size (radix on dest low bits) so static is well balanced, and
+                // bucket order does not affect the graph.
+#pragma omp for schedule(static)
+                for (int64_t b = 0; b < static_cast<int64_t>(kNumBuckets);
+                     b++) {
+                    Edge* p = base + bstart[b];
+                    size_t m = bstart[b + 1] - bstart[b];
+                    if (m == 0) {
+                        continue;
+                    }
+                    // Order within the bucket by (level, dest); no src
+                    // tiebreak is needed because the merge re-sorts each
+                    // node's incoming set.
+                    std::sort(p, p + m, [](const Edge& a, const Edge& c) {
+                        if (a.level != c.level) {
+                            return a.level < c.level;
+                        }
+                        return a.dest < c.dest;
+                    });
+                    std::vector<HNSW::storage_idx_t> incoming;
+                    size_t g = 0;
+                    while (g < m) {
+                        size_t h = g + 1;
+                        while (h < m && p[h].level == p[g].level &&
+                               p[h].dest == p[g].dest) {
+                            h++;
+                        }
+                        incoming.clear();
+                        incoming.reserve(h - g);
+                        for (size_t kk = g; kk < h; kk++) {
+                            incoming.push_back(p[kk].src);
+                        }
+                        hnsw.merge_reverse_links_deterministic(
+                                *dis,
+                                p[g].dest,
+                                p[g].level,
+                                incoming,
+                                keep_max_size_level0 && (p[g].level == 0));
+                        g = h;
+                    }
+                }
+            }
+
+            InterruptCallback::check();
+            s = e;
         }
+
+        i1 = i0;
     }
+
     if (verbose) {
         printf("Done in %.3f ms\n", getmillisecs() - t0);
     }
-    if (!index_hnsw.retain_locks) {
-        locks.clear();
-    }
 }
 
 } // namespace
@@ -384,13 +561,8 @@ void IndexHNSW::add(idx_t n, const float* x) {
     storage->add(n, x);
     ntotal = storage->ntotal;
 
-    hnsw_add_vertices(
-            *this,
-            n0,
-            n,
-            x,
-            verbose,
-            hnsw.levels.size() == static_cast<size_t>(ntotal));
+    bool preset_levels = hnsw.levels.size() == static_cast<size_t>(ntotal);
+    hnsw_add_vertices_deterministic(*this, n0, n, x, verbose, preset_levels);
 }
 
 void IndexHNSW::reset() {
 
@@ -53,7 +53,8 @@ struct IndexHNSW : Index {
     // See impl/VisitedTable.h.
     std::optional<bool> use_visited_hashset;
 
-    // Per-node locks for HNSW graph construction.
+    // Per-node locks, retained for the dormant init_level_0_from_entry_points
+    // helper and for IndexBinaryHNSW. The primary add() path is lock-free.
     LockVector locks;
     // locks are freed after each call to add() unless this flag is set.
     bool retain_locks = false;