faiss 0.6.2 → 0.6.4

data/vendor/faiss/faiss/IndexHNSW.cpp

@@ -7,7 +7,8 @@
 
 #include <faiss/IndexHNSW.h>
 
-#include <omp.h>
+#include <faiss/IndexRaBitQ.h>
+
 #include <atomic>
 #include <cinttypes>
 #include <cstdio>
@@ -17,7 +18,6 @@
 #include <limits>
 #include <memory>
 #include <queue>
-#include <random>
 
 #include <cstdint>
 #include "faiss/Index.h"
@@ -30,6 +30,7 @@
 #include <faiss/impl/FaissException.h>
 #include <faiss/impl/ResultHandler.h>
 #include <faiss/impl/VisitedTable.h>
+#include <faiss/impl/hnsw/LockVector.h>
 #include <faiss/impl/hnsw/MinimaxHeap.h>
 #include <faiss/utils/random.h>
 #include <faiss/utils/sorting.h>
@@ -60,19 +61,29 @@ DistanceComputer* storage_distance_computer(const Index* storage) {
     return storage->get_distance_computer();
 }
 
-void hnsw_add_vertices(
-        IndexHNSW& index_hnsw,
+} // namespace
+
+// Deterministic HNSW graph build inspired by ParlayANN: points are bucketed by
+// level and inserted in prefix-doubling batches. Phase A computes forward
+// links against a frozen snapshot and records reverse edges; phase B merges
+// those in a fixed order.
+// https://arxiv.org/abs/2305.04359
+void hnsw_add_vertices_deterministic(
+        HNSW& hnsw,
         size_t n0,
         size_t n,
-        const float* x,
+        int d,
+        bool init_level0,
+        bool keep_max_size_level0,
+        bool preset_levels,
         bool verbose,
-        bool preset_levels = false) {
-    size_t d = index_hnsw.d;
-    HNSW& hnsw = index_hnsw.hnsw;
+        const std::function<DistanceComputer*()>& make_distance_computer,
+        const std::function<void(DistanceComputer&, HNSW::storage_idx_t)>&
+                set_query) {
     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 +94,16 @@ void hnsw_add_vertices(
         return;
     }
 
-    int max_level = hnsw.prepare_level_tab(n, preset_levels);
+    // init_level0=false (CAGRA import) skips the level-0-only bucket; higher-
+    // level points still build their own level-0 links.
+    const int min_bucket_level = init_level0 ? 0 : 1;
 
+    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.
     std::vector<int> hist;
     std::vector<int> order(n);
 
@@ -122,98 +133,243 @@ void hnsw_add_vertices(
         }
     }
 
-    idx_t check_period = InterruptCallback::get_period_hint(
-            max_level * index_hnsw.d * hnsw.efConstruction);
+    // 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));
 
-    { // perform add
-        RandomGenerator rng2(789);
+    // Polled inside phase A: a batch can be 2% of ntotal, so cancelling only
+    // at batch boundaries would take minutes on large indexes.
+    const idx_t check_period = InterruptCallback::get_period_hint(
+            max_level * d * hnsw.efConstruction);
 
-        size_t i1 = static_cast<int>(n);
+    RandomGenerator rng2(789);
+    size_t i1 = n;
 
-        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];
+    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;
+        }
 
-            if (verbose) {
-                printf("Adding %zu elements at level %d\n", i1 - i0, pt_level);
-            }
+        if (verbose) {
+            printf("Adding %zu elements at level %d\n", i1 - i0, 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))]);
+        // 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))]);
+        }
+
+        // Bootstrap only when the graph is empty. Raising entry_point to a
+        // point that is not yet linked orphans everything already inserted,
+        // which on an incremental add() is the entire prior graph. order[i0]
+        // is the sole member of the first batch below, so defer until then.
+        bool raise_entry_point = false;
+        if (hnsw.entry_point == -1) {
+            hnsw.max_level = pt_level;
+            hnsw.entry_point = order[i0];
+        } else if (pt_level > hnsw.max_level) {
+            raise_entry_point = true;
+        }
+
+        // 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
+            // A reverse edge: `dest` gains an incoming link from `src`.
+            struct Edge {
+                int level;
+                HNSW::storage_idx_t dest;
+                HNSW::storage_idx_t src;
+            };
+
+            // One buffer, sized to an upper bound so it never reallocates.
+            // The fill order is nondeterministic but harmless: phase B
+            // regroups by destination.
+            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];
             }
+            // Default-initialised to skip a memset per sub-batch; safe because
+            // edge_counter only hands out slots that are written before read.
+            std::unique_ptr<Edge[]> reverse_edges(new Edge[cap_sum]);
+            std::atomic<size_t> edge_counter{0};
 
+            // Thrown after the region (cannot throw out of one); the
+            // unsynchronized write costs at most one extra poll.
             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::unique_ptr<DistanceComputer> dis(make_distance_computer());
+                std::vector<std::pair<HNSW::storage_idx_t, int>>
+                        pt_reverse_edges;
                 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;
+                    set_query(*dis, pt_id);
 
-                    hnsw.add_with_locks(
+                    pt_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);
+                            pt_reverse_edges,
+                            keep_max_size_level0);
+
+                    size_t off = edge_counter.fetch_add(
+                            pt_reverse_edges.size(), std::memory_order_relaxed);
+                    for (size_t k = 0; k < pt_reverse_edges.size(); k++) {
+                        reverse_edges[off + k] = {
+                                pt_reverse_edges[k].second,
+                                pt_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: merge the reverse edges
+            // Group by destination so each node is merged by exactly one
+            // thread: lock-free and thread-count-independent.
+            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>(reverse_edges[idx].dest) &
+                        kBucketMask) +
+                       1]++;
+            }
+            for (uint32_t b = 0; b < kNumBuckets; b++) {
+                bstart[b + 1] += bstart[b];
+            }
+
+            // In-place partition (cycle sort): each step lands at least one
+            // edge in its bucket, so O(total) with no auxiliary 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 = reverse_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, reverse_edges[head[tb]]);
+                            head[tb]++;
+                        }
+                        reverse_edges[head[b]] = cur_e;
+                        head[b]++;
+                    }
+                }
+            }
+            Edge* base = reverse_edges.get();
+
+#pragma omp parallel if (total > 100)
+            {
+                std::unique_ptr<DistanceComputer> dis(make_distance_computer());
+                // Not schedule(dynamic): the libomp dynamic dispatcher
+                // segfaults in some build configs.
+#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;
+                    }
+                    // src order does not matter; the merge re-sorts each 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;
+
+            if (raise_entry_point) {
+                hnsw.max_level = pt_level;
+                hnsw.entry_point = order[i0];
+                raise_entry_point = false;
+            }
         }
+
+        i1 = i0;
     }
+
+    if (init_level0) {
+        FAISS_ASSERT(i1 == 0);
+    } else {
+        FAISS_ASSERT((i1 - hist[0]) == 0);
+    }
+
     if (verbose) {
         printf("Done in %.3f ms\n", getmillisecs() - t0);
     }
-    if (!index_hnsw.retain_locks) {
-        locks.clear();
-    }
 }
 
-} // namespace
-
 /**************************************************************
  * IndexHNSW implementation
  **************************************************************/
@@ -269,6 +425,7 @@ void hnsw_search(
         }
     }
     size_t n1 = 0, n2 = 0, ndis = 0, nhops = 0;
+    size_t n_rabitq_1bit = 0, n_rabitq_refine = 0;
 
     idx_t check_period = InterruptCallback::get_period_hint(
             hnsw.max_level * index->d * efSearch);
@@ -280,12 +437,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);
@@ -294,7 +451,9 @@ void hnsw_search(
                 omp_capture_exception(ex, [&] { interrupt = true; });
             }
 
-#pragma omp for reduction(+ : n1, n2, ndis, nhops) schedule(guided)
+#pragma omp for reduction(                                               \
+                + : n1, n2, ndis, nhops, n_rabitq_1bit, n_rabitq_refine) \
+        schedule(guided)
             for (idx_t i = i0; i < i1; i++) {
                 if (interrupt.load(std::memory_order_relaxed)) {
                     continue;
@@ -302,6 +461,10 @@ void hnsw_search(
                 try {
                     res->begin(i);
                     dis->set_query(x + i * index->d);
+                    auto* rq = dynamic_cast<RaBitQDistanceComputer*>(dis.get());
+                    if (rq) {
+                        rq->stats.reset();
+                    }
 
                     HNSWStats stats =
                             hnsw.search(*dis, index, *res, *vt, params);
@@ -309,6 +472,10 @@ void hnsw_search(
                     n2 += stats.n2;
                     ndis += stats.ndis;
                     nhops += stats.nhops;
+                    if (rq) {
+                        n_rabitq_1bit += rq->stats.n_1bit;
+                        n_rabitq_refine += rq->stats.n_refine;
+                    }
                     res->end();
                     vt->advance();
                 } catch (...) {
@@ -321,6 +488,7 @@ void hnsw_search(
     }
 
     hnsw_stats.combine({n1, n2, ndis, nhops});
+    rabitq_stats.add({n_rabitq_1bit, n_rabitq_refine});
 }
 
 } // anonymous namespace
@@ -384,18 +552,25 @@ void IndexHNSW::add(idx_t n, const float* x) {
     storage->add(n, x);
     ntotal = storage->ntotal;
 
-    hnsw_add_vertices(
-            *this,
+    bool preset_levels = hnsw.levels.size() == static_cast<size_t>(ntotal);
+
+    hnsw_add_vertices_deterministic(
+            hnsw,
             n0,
             n,
-            x,
+            d,
+            init_level0,
+            keep_max_size_level0,
+            preset_levels,
             verbose,
-            hnsw.levels.size() == static_cast<size_t>(ntotal));
+            [this] { return storage_distance_computer(storage); },
+            [this, x, n0](DistanceComputer& dc, HNSW::storage_idx_t pt_id) {
+                dc.set_query(x + (pt_id - n0) * d);
+            });
 }
 
 void IndexHNSW::reset() {
     hnsw.reset();
-    locks.clear();
     storage->reset();
     ntotal = 0;
 }
@@ -479,11 +654,12 @@ void IndexHNSW::search_level_0(
         {
             std::unique_ptr<DistanceComputer> qdis;
             HNSWStats search_stats;
-            std::unique_ptr<VisitedTable> vt;
+            RaBitQStats rq_search_stats;
+            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 (...) {
@@ -498,6 +674,11 @@ void IndexHNSW::search_level_0(
                 try {
                     res->begin(i);
                     qdis->set_query(x + i * d);
+                    auto* rq =
+                            dynamic_cast<RaBitQDistanceComputer*>(qdis.get());
+                    if (rq) {
+                        rq->stats.reset();
+                    }
 
                     hnsw.search_level_0(
                             *qdis.get(),
@@ -509,6 +690,9 @@ void IndexHNSW::search_level_0(
                             search_stats,
                             *vt,
                             params);
+                    if (rq) {
+                        rq_search_stats.add(rq->stats);
+                    }
                     res->end();
                     vt->advance();
                 } catch (...) {
@@ -518,6 +702,7 @@ void IndexHNSW::search_level_0(
 #pragma omp critical
             {
                 hnsw_stats.combine(search_stats);
+                rabitq_stats.add(rq_search_stats);
             }
         }
         omp_rethrow_if_exception(ex);
@@ -576,7 +761,7 @@ void IndexHNSW::init_level_0_from_entry_points(
         int n,
         const storage_idx_t* points,
         const storage_idx_t* nearests) {
-    locks.prepare(ntotal);
+    LockVector locks(ntotal);
 
 #pragma omp parallel
     {
@@ -606,10 +791,6 @@ void IndexHNSW::init_level_0_from_entry_points(
     if (verbose) {
         printf("\n");
     }
-
-    if (!retain_locks) {
-        locks.clear();
-    }
 }
 
 void IndexHNSW::reorder_links() {
@@ -741,7 +922,7 @@ IndexHNSWFlatPanorama::IndexHNSWFlatPanorama(
     // Enable Panorama search mode.
     // This is not ideal, but is still more simple than making a subclass of
     // HNSW and overriding the search logic.
-    hnsw.is_panorama = true;
+    hnsw.search_method = HNSW::SM_PANORAMA;
 }
 
 void IndexHNSWFlatPanorama::add(idx_t n, const float* x) {
@@ -808,6 +989,38 @@ IndexHNSWSQ::IndexHNSWSQ(
 
 IndexHNSWSQ::IndexHNSWSQ() = default;
 
+/**************************************************************
+ * IndexHNSWRaBitQ implementation
+ **************************************************************/
+
+IndexHNSWRaBitQ::IndexHNSWRaBitQ() = default;
+
+namespace {
+
+IndexRaBitQ* make_hnsw_rabitq_storage(
+        int d,
+        uint8_t nb_bits,
+        MetricType metric) {
+    FAISS_THROW_IF_NOT_MSG(
+            metric == METRIC_L2, "IndexHNSWRaBitQ supports only the L2 metric");
+    return new IndexRaBitQ(d, metric, nb_bits);
+}
+
+} // namespace
+
+IndexHNSWRaBitQ::IndexHNSWRaBitQ(
+        int d,
+        int M,
+        uint8_t nb_bits,
+        MetricType metric)
+        : IndexHNSW(make_hnsw_rabitq_storage(d, nb_bits, metric), M) {
+    own_fields = true;
+    is_trained = storage->is_trained;
+    // 1-bit codes store plain SignBitFactors with no f_error, so there is no
+    // bound to prune with and the staged path does not apply.
+    hnsw.search_method = nb_bits >= 2 ? HNSW::SM_RABITQ : HNSW::SM_DEFAULT;
+}
+
 /**************************************************************
  * IndexHNSW2Level implementation
  **************************************************************/
@@ -903,8 +1116,7 @@ void IndexHNSW2Level::search(
         idx_t* labels,
         const SearchParameters* params) const {
     FAISS_THROW_IF_NOT(k > 0);
-    FAISS_THROW_IF_NOT_MSG(
-            !params, "search params not supported for this index");
+    FAISS_THROW_IF_MSG(params, "search params not supported for this index");
 
     if (dynamic_cast<const Index2Layer*>(storage)) {
         IndexHNSW::search(n, x, k, distances, labels);
@@ -1084,8 +1296,8 @@ IndexHNSWCagra::IndexHNSWCagra(
 }
 
 void IndexHNSWCagra::add(idx_t n, const float* x) {
-    FAISS_THROW_IF_NOT_MSG(
-            !base_level_only,
+    FAISS_THROW_IF_MSG(
+            base_level_only,
             "Cannot add vectors when base_level_only is set to True");
 
     IndexHNSW::add(n, x);
@@ -1124,13 +1336,11 @@ void IndexHNSWCagra::search(
                 // first real candidate will always be strictly better.
                 nearest_d[i] = C::neutral();
 
-                std::random_device rd;
-                std::mt19937 gen(rd());
-                std::uniform_int_distribution<idx_t> distrib(
-                        0, this->ntotal - 1);
+                // Seeded per query so entrypoints are reproducible.
+                SplitMix64RandomGenerator gen(i);
 
                 for (idx_t j = 0; j < num_base_level_search_entrypoints; j++) {
-                    auto idx = distrib(gen);
+                    idx_t idx = gen.rand_int64() % this->ntotal;
                     auto distance = (*dis)(idx);
                     if (C::cmp(nearest_d[i], distance)) {
                         nearest[i] = static_cast<storage_idx_t>(idx);
@@ -1189,12 +1399,11 @@ void IndexHNSWCagra::range_search(
             // real candidate will always be strictly better.
             float nearest_d = C::neutral();
 
-            std::random_device rd;
-            std::mt19937 gen(rd());
-            std::uniform_int_distribution<idx_t> distrib(0, ntotal - 1);
+            // For reproducible entrypoint.
+            SplitMix64RandomGenerator gen(i);
 
             for (idx_t j = 0; j < num_base_level_search_entrypoints; j++) {
-                auto idx = distrib(gen);
+                idx_t idx = gen.rand_int64() % ntotal;
                 auto distance = (*dis)(idx);
                 // C::cmp(nearest_d, distance) is true iff distance is
                 // strictly better than the current nearest_d.
@@ -1208,11 +1417,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;