faiss 0.5.3 → 0.6.1

data/vendor/faiss/faiss/IndexBinaryHNSW.cpp

@@ -22,12 +22,23 @@
 #include <faiss/impl/DistanceComputer.h>
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/ResultHandler.h>
+#include <faiss/impl/VisitedTable.h>
 #include <faiss/utils/Heap.h>
 #include <faiss/utils/hamming.h>
 #include <faiss/utils/random.h>
 
+#include <faiss/impl/simd_dispatch.h>
+
 #include <random>
 
+// Scalar (NONE) fallback for dynamic dispatch
+#define THE_SIMD_LEVEL SIMDLevel::NONE
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/impl/binary_hamming/IndexBinaryHNSW_impl.h>
+#include <faiss/utils/hamming_distance/hamming_computer-generic.h>
+#undef THE_SIMD_LEVEL
+
 namespace faiss {
 
 /**************************************************************
@@ -60,10 +71,8 @@ void hnsw_add_vertices(
         printf("  max_level = %d\n", max_level);
     }
 
-    std::vector<omp_lock_t> locks(ntotal);
-    for (int i = 0; i < ntotal; i++) {
-        omp_init_lock(&locks[i]);
-    }
+    auto& locks = index_hnsw.locks;
+    locks.prepare(ntotal);
 
     // add vectors from highest to lowest level
     std::vector<int> hist;
@@ -72,10 +81,11 @@ void hnsw_add_vertices(
     { // make buckets with vectors of the same level
 
         // build histogram
-        for (int i = 0; i < n; i++) {
-            HNSW::storage_idx_t pt_id = i + n0;
+        for (size_t i = 0; i < n; i++) {
+            HNSW::storage_idx_t pt_id =
+                    static_cast<HNSW::storage_idx_t>(i + n0);
             int pt_level = hnsw.levels[pt_id] - 1;
-            while (pt_level >= hist.size()) {
+            while (pt_level >= static_cast<int>(hist.size())) {
                 hist.push_back(0);
             }
             hist[pt_level]++;
@@ -83,13 +93,14 @@ void hnsw_add_vertices(
 
         // accumulate
         std::vector<int> offsets(hist.size() + 1, 0);
-        for (int i = 0; i < hist.size() - 1; i++) {
+        for (size_t i = 0; i < hist.size() - 1; i++) {
             offsets[i + 1] = offsets[i] + hist[i];
         }
 
         // bucket sort
-        for (int i = 0; i < n; i++) {
-            HNSW::storage_idx_t pt_id = i + n0;
+        for (size_t i = 0; i < n; i++) {
+            HNSW::storage_idx_t pt_id =
+                    static_cast<HNSW::storage_idx_t>(i + n0);
             int pt_level = hnsw.levels[pt_id] - 1;
             order[offsets[pt_level]++] = pt_id;
         }
@@ -98,20 +109,22 @@ void hnsw_add_vertices(
     { // perform add
         RandomGenerator rng2(789);
 
-        int i1 = n;
+        size_t i1 = static_cast<int>(n);
 
-        for (int pt_level = hist.size() - 1;
+        for (int pt_level = static_cast<int>(hist.size()) - 1;
              pt_level >= int(!index_hnsw.init_level0);
              pt_level--) {
-            int i0 = i1 - hist[pt_level];
+            size_t i0 = i1 - hist[pt_level];
 
             if (verbose) {
-                printf("Adding %d elements at level %d\n", i1 - i0, pt_level);
+                printf("Adding %zu elements at level %d\n", i1 - i0, pt_level);
             }
 
             // random permutation to get rid of dataset order bias
-            for (int j = i0; j < i1; j++) {
-                std::swap(order[j], order[j + rng2.rand_int(i1 - j)]);
+            for (size_t j = i0; j < i1; j++) {
+                std::swap(
+                        order[j],
+                        order[j + rng2.rand_int(static_cast<int>(i1 - j))]);
             }
 
 #pragma omp parallel
@@ -120,11 +133,11 @@ void hnsw_add_vertices(
 
                 std::unique_ptr<DistanceComputer> dis(
                         index_hnsw.get_distance_computer());
-                int prev_display =
-                        verbose && omp_get_thread_num() == 0 ? 0 : -1;
+                bool do_display = verbose && omp_get_thread_num() == 0;
+                size_t prev_display = 0;
 
 #pragma omp for schedule(dynamic)
-                for (int i = i0; i < i1; i++) {
+                for (int64_t i = i0; i < i1; i++) {
                     HNSW::storage_idx_t pt_id = order[i];
                     dis->set_query(
                             (float*)(x + (pt_id - n0) * index_hnsw.code_size));
@@ -137,9 +150,9 @@ void hnsw_add_vertices(
                             vt,
                             index_hnsw.keep_max_size_level0 && (pt_level == 0));
 
-                    if (prev_display >= 0 && i - i0 > prev_display + 10000) {
+                    if (do_display && i - i0 > prev_display + 10000) {
                         prev_display = i - i0;
-                        printf("  %d / %d\r", i - i0, i1 - i0);
+                        printf("  %zu / %zu\r", i - i0, i1 - i0);
                         fflush(stdout);
                     }
                 }
@@ -155,9 +168,8 @@ void hnsw_add_vertices(
     if (verbose) {
         printf("Done in %.3f ms\n", getmillisecs() - t0);
     }
-
-    for (int i = 0; i < ntotal; i++) {
-        omp_destroy_lock(&locks[i]);
+    if (!index_hnsw.retain_locks) {
+        locks.clear();
     }
 }
 
@@ -171,19 +183,19 @@ IndexBinaryHNSW::IndexBinaryHNSW() {
     is_trained = true;
 }
 
-IndexBinaryHNSW::IndexBinaryHNSW(int d, int M)
-        : IndexBinary(d),
+IndexBinaryHNSW::IndexBinaryHNSW(int d_, int M)
+        : IndexBinary(d_),
           hnsw(M),
           own_fields(true),
-          storage(new IndexBinaryFlat(d)) {
+          storage(new IndexBinaryFlat(d_)) {
     is_trained = true;
 }
 
-IndexBinaryHNSW::IndexBinaryHNSW(IndexBinary* storage, int M)
-        : IndexBinary(storage->d),
+IndexBinaryHNSW::IndexBinaryHNSW(IndexBinary* storage_, int M)
+        : IndexBinary(storage_->d),
           hnsw(M),
           own_fields(false),
-          storage(storage) {
+          storage(storage_) {
     is_trained = true;
 }
 
@@ -205,10 +217,14 @@ void IndexBinaryHNSW::search(
         idx_t k,
         int32_t* distances,
         idx_t* labels,
-        const SearchParameters* params) const {
-    FAISS_THROW_IF_NOT_MSG(
-            !params, "search params not supported for this index");
+        const SearchParameters* params_in) const {
     FAISS_THROW_IF_NOT(k > 0);
+    const SearchParametersHNSW* params = nullptr;
+    if (params_in) {
+        params = dynamic_cast<const SearchParametersHNSW*>(params_in);
+        FAISS_THROW_IF_NOT_MSG(
+                params, "IndexBinaryHNSW params have incorrect type");
+    }
 
     // we use the buffer for distances as float but convert them back
     // to int in the end
@@ -217,42 +233,57 @@ void IndexBinaryHNSW::search(
     using RH = HeapBlockResultHandler<HNSW::C>;
     RH bres(n, distances_f, labels, k);
 
+    size_t n1 = 0, n2 = 0, ndis = 0, nhops = 0;
+
 #pragma omp parallel
     {
         VisitedTable vt(ntotal);
         std::unique_ptr<DistanceComputer> dis(get_distance_computer());
         RH::SingleResultHandler res(bres);
 
-#pragma omp for
+#pragma omp for reduction(+ : n1, n2, ndis, nhops)
         for (idx_t i = 0; i < n; i++) {
             res.begin(i);
             dis->set_query((float*)(x + i * code_size));
             // Given that IndexBinaryHNSW is not an IndexHNSW, we pass nullptr
             // as the index parameter. This state does not get used in the
-            // search function, as it is merely there to to enable Panorama
+            // search function, as it is merely there to enable Panorama
             // execution for IndexHNSWFlatPanorama.
-            hnsw.search(*dis, nullptr, res, vt);
+            HNSWStats stats = hnsw.search(*dis, nullptr, res, vt, params_in);
+            n1 += stats.n1;
+            n2 += stats.n2;
+            ndis += stats.ndis;
+            nhops += stats.nhops;
             res.end();
         }
     }
 
+    hnsw_stats.combine({n1, n2, ndis, nhops});
+
 #pragma omp parallel for
-    for (int i = 0; i < n * k; ++i) {
+    for (idx_t i = 0; i < n * k; ++i) {
         distances[i] = std::round(distances_f[i]);
     }
 }
 
 void IndexBinaryHNSW::add(idx_t n, const uint8_t* x) {
     FAISS_THROW_IF_NOT(is_trained);
-    int n0 = ntotal;
+    size_t n0 = ntotal;
     storage->add(n, x);
     ntotal = storage->ntotal;
 
-    hnsw_add_vertices(*this, n0, n, x, verbose, hnsw.levels.size() == ntotal);
+    hnsw_add_vertices(
+            *this,
+            n0,
+            n,
+            x,
+            verbose,
+            hnsw.levels.size() == static_cast<size_t>(ntotal));
 }
 
 void IndexBinaryHNSW::reset() {
     hnsw.reset();
+    locks.clear();
     storage->reset();
     ntotal = 0;
 }
@@ -261,60 +292,15 @@ void IndexBinaryHNSW::reconstruct(idx_t key, uint8_t* recons) const {
     storage->reconstruct(key, recons);
 }
 
-namespace {
-
-template <class HammingComputer>
-struct FlatHammingDis : DistanceComputer {
-    const int code_size;
-    const uint8_t* b;
-    size_t ndis;
-    HammingComputer hc;
-
-    float operator()(idx_t i) override {
-        ndis++;
-        return hc.hamming(b + i * code_size);
-    }
-
-    float symmetric_dis(idx_t i, idx_t j) override {
-        return HammingComputerDefault(b + j * code_size, code_size)
-                .hamming(b + i * code_size);
-    }
-
-    explicit FlatHammingDis(const IndexBinaryFlat& storage)
-            : code_size(storage.code_size),
-              b(storage.xb.data()),
-              ndis(0),
-              hc() {}
-
-    // NOTE: Pointers are cast from float in order to reuse the floating-point
-    //   DistanceComputer.
-    void set_query(const float* x) override {
-        hc.set((uint8_t*)x, code_size);
-    }
-
-    ~FlatHammingDis() override {
-#pragma omp critical
-        {
-            hnsw_stats.ndis += ndis;
-        }
-    }
-};
-
-struct BuildDistanceComputer {
-    using T = DistanceComputer*;
-    template <class HammingComputer>
-    DistanceComputer* f(IndexBinaryFlat* flat_storage) {
-        return new FlatHammingDis<HammingComputer>(*flat_storage);
-    }
-};
-
-} // namespace
-
 DistanceComputer* IndexBinaryHNSW::get_distance_computer() const {
     IndexBinaryFlat* flat_storage = dynamic_cast<IndexBinaryFlat*>(storage);
-    FAISS_ASSERT(flat_storage != nullptr);
-    BuildDistanceComputer bd;
-    return dispatch_HammingComputer(code_size, bd, flat_storage);
+    FAISS_THROW_IF_NOT_MSG(
+            flat_storage != nullptr,
+            "IndexBinaryHNSW requires IndexBinaryFlat storage");
+    return with_simd_level([&]<SIMDLevel SL>() {
+        return make_binary_hnsw_distance_computer_fixSL<SL>(
+                code_size, flat_storage);
+    });
 }
 
 /**************************************************************
@@ -325,8 +311,8 @@ IndexBinaryHNSWCagra::IndexBinaryHNSWCagra() : IndexBinaryHNSW() {
     storage = nullptr;
 }
 
-IndexBinaryHNSWCagra::IndexBinaryHNSWCagra(int d, int M)
-        : IndexBinaryHNSW(d, M) {
+IndexBinaryHNSWCagra::IndexBinaryHNSWCagra(int d_, int M)
+        : IndexBinaryHNSW(d_, M) {
     init_level0 = true;
     keep_max_size_level0 = true;
 }
@@ -349,6 +335,13 @@ void IndexBinaryHNSWCagra::search(
     if (!base_level_only) {
         IndexBinaryHNSW::search(n, x, k, distances, labels, params);
     } else {
+        FAISS_THROW_IF_NOT_MSG(
+                ntotal > 0, "IndexBinaryHNSWCagra: cannot search empty index");
+        FAISS_THROW_IF_NOT_MSG(
+                num_base_level_search_entrypoints > 0,
+                "IndexBinaryHNSWCagra: "
+                "num_base_level_search_entrypoints must be > 0");
+
         float* distances_f = (float*)distances;
 
         using RH = HeapBlockResultHandler<HNSW::C>;
@@ -374,7 +367,7 @@ void IndexBinaryHNSWCagra::search(
                 float distance = (*dis)(idx);
 
                 if (distance < nearest_d[i]) {
-                    nearest[i] = idx;
+                    nearest[i] = static_cast<storage_idx_t>(idx);
                     nearest_d[i] = distance;
                 }
             }
@@ -407,10 +400,14 @@ void IndexBinaryHNSWCagra::search(
 
                 res.end();
             }
+#pragma omp critical
+            {
+                hnsw_stats.combine(search_stats);
+            }
         }
 
 #pragma omp parallel for
-        for (int i = 0; i < n * k; ++i) {
+        for (idx_t i = 0; i < n * k; ++i) {
             distances[i] = std::round(distances_f[i]);
         }
     }

data/vendor/faiss/faiss/IndexBinaryHNSW.h

@@ -11,6 +11,7 @@
 
 #include <faiss/IndexBinaryFlat.h>
 #include <faiss/impl/HNSW.h>
+#include <faiss/impl/hnsw/LockVector.h>
 #include <faiss/utils/utils.h>
 
 namespace faiss {
@@ -19,14 +20,14 @@ namespace faiss {
  * link structure built on top */
 
 struct IndexBinaryHNSW : IndexBinary {
-    typedef HNSW::storage_idx_t storage_idx_t;
+    using storage_idx_t = HNSW::storage_idx_t;
 
     // the link structure
     HNSW hnsw;
 
     // the sequential storage
-    bool own_fields;
-    IndexBinary* storage;
+    bool own_fields = false;
+    IndexBinary* storage = nullptr;
 
     // When set to false, level 0 in the knn graph is not initialized.
     // This option is used by GpuIndexBinaryCagra::copyTo(IndexBinaryHNSW*)
@@ -40,6 +41,11 @@ struct IndexBinaryHNSW : IndexBinary {
     // used when GpuIndexBinaryCagra::copyFrom(IndexBinaryHNSW*) is called.
     bool keep_max_size_level0 = false;
 
+    // Per-node locks for HNSW graph construction.
+    LockVector locks;
+    // locks are freed after each call to add() unless this flag is set.
+    bool retain_locks = false;
+
     explicit IndexBinaryHNSW();
     explicit IndexBinaryHNSW(int d, int M = 32);
     explicit IndexBinaryHNSW(IndexBinary* storage, int M = 32);