faiss 0.4.1 → 0.4.2

data/vendor/faiss/faiss/IndexBinaryHNSW.cpp

@@ -26,6 +26,8 @@
 #include <faiss/utils/hamming.h>
 #include <faiss/utils/random.h>
 
+#include <random>
+
 namespace faiss {
 
 /**************************************************************
@@ -98,7 +100,9 @@ void hnsw_add_vertices(
 
         int i1 = n;
 
-        for (int pt_level = hist.size() - 1; pt_level >= 0; pt_level--) {
+        for (int pt_level = hist.size() - 1;
+             pt_level >= int(!index_hnsw.init_level0);
+             pt_level--) {
             int i0 = i1 - hist[pt_level];
 
             if (verbose) {
@@ -125,7 +129,13 @@ void hnsw_add_vertices(
                     dis->set_query(
                             (float*)(x + (pt_id - n0) * index_hnsw.code_size));
 
-                    hnsw.add_with_locks(*dis, pt_level, pt_id, locks, vt);
+                    hnsw.add_with_locks(
+                            *dis,
+                            pt_level,
+                            pt_id,
+                            locks,
+                            vt,
+                            index_hnsw.keep_max_size_level0 && (pt_level == 0));
 
                     if (prev_display >= 0 && i - i0 > prev_display + 10000) {
                         prev_display = i - i0;
@@ -136,14 +146,19 @@ void hnsw_add_vertices(
             }
             i1 = i0;
         }
-        FAISS_ASSERT(i1 == 0);
+        if (index_hnsw.init_level0) {
+            FAISS_ASSERT(i1 == 0);
+        } else {
+            FAISS_ASSERT((i1 - hist[0]) == 0);
+        }
     }
     if (verbose) {
         printf("Done in %.3f ms\n", getmillisecs() - t0);
     }
 
-    for (int i = 0; i < ntotal; i++)
+    for (int i = 0; i < ntotal; i++) {
         omp_destroy_lock(&locks[i]);
+    }
 }
 
 } // anonymous namespace
@@ -296,4 +311,103 @@ DistanceComputer* IndexBinaryHNSW::get_distance_computer() const {
     return dispatch_HammingComputer(code_size, bd, flat_storage);
 }
 
+/**************************************************************
+ * IndexBinaryHNSWCagra implementation
+ **************************************************************/
+
+IndexBinaryHNSWCagra::IndexBinaryHNSWCagra() : IndexBinaryHNSW() {
+    storage = nullptr;
+}
+
+IndexBinaryHNSWCagra::IndexBinaryHNSWCagra(int d, int M)
+        : IndexBinaryHNSW(d, M) {
+    init_level0 = true;
+    keep_max_size_level0 = true;
+}
+
+void IndexBinaryHNSWCagra::add(idx_t n, const uint8_t* x) {
+    FAISS_THROW_IF_NOT_MSG(
+            !base_level_only,
+            "Cannot add vectors when base_level_only is set to True");
+
+    IndexBinaryHNSW::add(n, x);
+}
+
+void IndexBinaryHNSWCagra::search(
+        idx_t n,
+        const uint8_t* x,
+        idx_t k,
+        int32_t* distances,
+        idx_t* labels,
+        const SearchParameters* params) const {
+    if (!base_level_only) {
+        IndexBinaryHNSW::search(n, x, k, distances, labels, params);
+    } else {
+        float* distances_f = (float*)distances;
+
+        using RH = HeapBlockResultHandler<HNSW::C>;
+        RH bres(n, distances_f, labels, k);
+
+        std::vector<storage_idx_t> nearest(n);
+        std::vector<float> nearest_d(n);
+
+#pragma omp parallel for
+        for (idx_t i = 0; i < n; i++) {
+            std::unique_ptr<DistanceComputer> dis(get_distance_computer());
+            dis->set_query((float*)(x + i * code_size));
+
+            nearest[i] = -1;
+            nearest_d[i] = std::numeric_limits<float>::max();
+
+            std::random_device rd;
+            std::mt19937 gen(rd());
+            std::uniform_int_distribution<idx_t> distrib(0, this->ntotal - 1);
+
+            for (idx_t j = 0; j < num_base_level_search_entrypoints; j++) {
+                auto idx = distrib(gen);
+                float distance = (*dis)(idx);
+
+                if (distance < nearest_d[i]) {
+                    nearest[i] = idx;
+                    nearest_d[i] = distance;
+                }
+            }
+            FAISS_THROW_IF_NOT_MSG(
+                    nearest[i] >= 0, "Could not find a valid entrypoint.");
+        }
+
+#pragma omp parallel
+        {
+            VisitedTable vt(ntotal);
+            std::unique_ptr<DistanceComputer> dis(get_distance_computer());
+            HNSWStats search_stats;
+            RH::SingleResultHandler res(bres);
+
+#pragma omp for
+            for (idx_t i = 0; i < n; i++) {
+                res.begin(i);
+                dis->set_query((float*)(x + i * code_size));
+
+                hnsw.search_level_0(
+                        *dis,
+                        res,
+                        1,
+                        &nearest[i],
+                        &nearest_d[i],
+                        1, // search_type
+                        search_stats,
+                        vt,
+                        params);
+
+                res.end();
+            }
+        }
+
+#pragma omp parallel for
+        for (int i = 0; i < n * k; ++i) {
+            distances[i] = std::round(distances_f[i]);
+        }
+    }
+}
+
 } // namespace faiss

data/vendor/faiss/faiss/IndexBinaryHNSW.h

@@ -28,6 +28,18 @@ struct IndexBinaryHNSW : IndexBinary {
     bool own_fields;
     IndexBinary* storage;
 
+    // When set to false, level 0 in the knn graph is not initialized.
+    // This option is used by GpuIndexBinaryCagra::copyTo(IndexBinaryHNSW*)
+    // as level 0 knn graph is copied over from the index built by
+    // GpuIndexBinaryCagra.
+    bool init_level0 = true;
+
+    // When set to true, all neighbors in level 0 are filled up
+    // to the maximum size allowed (2 * M). This option is used by
+    // IndexBinaryHHNSW to create a full base layer graph that is
+    // used when GpuIndexBinaryCagra::copyFrom(IndexBinaryHNSW*) is called.
+    bool keep_max_size_level0 = false;
+
     explicit IndexBinaryHNSW();
     explicit IndexBinaryHNSW(int d, int M = 32);
     explicit IndexBinaryHNSW(IndexBinary* storage, int M = 32);
@@ -55,4 +67,33 @@ struct IndexBinaryHNSW : IndexBinary {
     void reset() override;
 };
 
+struct IndexBinaryHNSWCagra : IndexBinaryHNSW {
+    IndexBinaryHNSWCagra();
+    IndexBinaryHNSWCagra(int d, int M);
+
+    /// When set to true, the index is immutable.
+    /// This option is used to copy the knn graph from GpuIndexBinaryCagra
+    /// to the base level of IndexBinaryHNSWCagra without adding upper levels.
+    /// Doing so enables to search the HNSW index, but removes the
+    /// ability to add vectors.
+    bool base_level_only = false;
+
+    /// When `base_level_only` is set to `True`, the search function
+    /// searches only the base level knn graph of the HNSW index.
+    /// This parameter selects the entry point by randomly selecting
+    /// some points and using the best one.
+    int num_base_level_search_entrypoints = 32;
+
+    void add(idx_t n, const uint8_t* x) override;
+
+    /// entry point for search
+    void search(
+            idx_t n,
+            const uint8_t* x,
+            idx_t k,
+            int32_t* distances,
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
+};
+
 } // namespace faiss

data/vendor/faiss/faiss/IndexBinaryHash.cpp

@@ -11,7 +11,6 @@
 
 #include <cinttypes>
 #include <cstdio>
-#include <memory>
 #include <unordered_set>
 
 #include <faiss/utils/hamming.h>

data/vendor/faiss/faiss/IndexBinaryIVF.cpp

@@ -26,9 +26,16 @@
 
 namespace faiss {
 
-IndexBinaryIVF::IndexBinaryIVF(IndexBinary* quantizer, size_t d, size_t nlist)
+IndexBinaryIVF::IndexBinaryIVF(
+        IndexBinary* quantizer,
+        size_t d,
+        size_t nlist,
+        bool own_invlists)
         : IndexBinary(d),
-          invlists(new ArrayInvertedLists(nlist, code_size)),
+          invlists(
+                  own_invlists ? new ArrayInvertedLists(nlist, code_size)
+                               : nullptr),
+          own_invlists(own_invlists),
           quantizer(quantizer),
           nlist(nlist) {
     FAISS_THROW_IF_NOT(d == quantizer->d);
@@ -283,7 +290,7 @@ void IndexBinaryIVF::check_compatible_for_merge(
             direct_map.no() && other->direct_map.no(),
             "direct map copy not implemented");
     FAISS_THROW_IF_NOT_MSG(
-            typeid(*this) == typeid(other),
+            typeid(*this) == typeid(*other),
             "can only merge indexes of the same type");
 }
 
@@ -444,8 +451,9 @@ void search_knn_hamming_heap(
                         list_size, scodes.get(), ids, simi, idxi, k);
 
                 nscan += list_size;
-                if (max_codes && nscan >= max_codes)
+                if (max_codes && nscan >= max_codes) {
                     break;
+                }
             }
 
             ndis += nscan;
@@ -532,8 +540,9 @@ void search_knn_hamming_count(
             }
 
             nscan += list_size;
-            if (max_codes && nscan >= max_codes)
+            if (max_codes && nscan >= max_codes) {
                 break;
+            }
         }
         ndis += nscan;
 
@@ -850,8 +859,9 @@ void IndexBinaryIVF::range_search_preassigned(
 
         auto scan_list_func = [&](size_t i, size_t ik, RangeQueryResult& qres) {
             idx_t key = assign[i * nprobe_2 + ik]; /* select the list  */
-            if (key < 0)
+            if (key < 0) {
                 return;
+            }
             FAISS_THROW_IF_NOT_FMT(
                     key < (idx_t)nlist,
                     "Invalid key=%" PRId64 " at ik=%zd nlist=%zd\n",
@@ -860,8 +870,9 @@ void IndexBinaryIVF::range_search_preassigned(
                     nlist);
             const size_t list_size = invlists->list_size(key);
 
-            if (list_size == 0)
+            if (list_size == 0) {
                 return;
+            }
 
             InvertedLists::ScopedCodes scodes(invlists, key);
             InvertedLists::ScopedIds ids(invlists, key);

data/vendor/faiss/faiss/IndexBinaryIVF.h

@@ -68,7 +68,11 @@ struct IndexBinaryIVF : IndexBinary {
      * identifier. The pointer is borrowed: the quantizer should not
      * be deleted while the IndexBinaryIVF is in use.
      */
-    IndexBinaryIVF(IndexBinary* quantizer, size_t d, size_t nlist);
+    IndexBinaryIVF(
+            IndexBinary* quantizer,
+            size_t d,
+            size_t nlist,
+            bool own_invlists = true);
 
     IndexBinaryIVF();
 

data/vendor/faiss/faiss/IndexFlat.cpp

@@ -239,6 +239,7 @@ FlatCodesDistanceComputer* IndexFlat::get_FlatCodesDistanceComputer() const {
 }
 
 void IndexFlat::reconstruct(idx_t key, float* recons) const {
+    FAISS_THROW_IF_NOT(key < ntotal);
     memcpy(recons, &(codes[key * code_size]), code_size);
 }
 
@@ -399,8 +400,9 @@ void IndexFlat1D::update_permutation() {
 
 void IndexFlat1D::add(idx_t n, const float* x) {
     IndexFlatL2::add(n, x);
-    if (continuous_update)
+    if (continuous_update) {
         update_permutation();
+    }
 }
 
 void IndexFlat1D::reset() {
@@ -452,10 +454,11 @@ void IndexFlat1D::search(
 
         while (i0 + 1 < i1) {
             idx_t imed = (i0 + i1) / 2;
-            if (xb[perm[imed]] <= q)
+            if (xb[perm[imed]] <= q) {
                 i0 = imed;
-            else
+            } else {
                 i1 = imed;
+            }
         }
 
         // query is between xb[perm[i0]] and xb[perm[i1]]
@@ -516,5 +519,4 @@ void IndexFlat1D::search(
     done:;
     }
 }
-
 } // namespace faiss

data/vendor/faiss/faiss/IndexHNSW.cpp

@@ -19,6 +19,7 @@
 #include <random>
 
 #include <cstdint>
+#include "faiss/Index.h"
 
 #include <faiss/Index2Layer.h>
 #include <faiss/IndexFlat.h>
@@ -81,8 +82,9 @@ void hnsw_add_vertices(
     }
 
     std::vector<omp_lock_t> locks(ntotal);
-    for (int i = 0; i < ntotal; i++)
+    for (int i = 0; i < ntotal; i++) {
         omp_init_lock(&locks[i]);
+    }
 
     // add vectors from highest to lowest level
     std::vector<int> hist;
@@ -94,8 +96,9 @@ void hnsw_add_vertices(
         for (int i = 0; i < n; i++) {
             storage_idx_t pt_id = i + n0;
             int pt_level = hnsw.levels[pt_id] - 1;
-            while (pt_level >= hist.size())
+            while (pt_level >= hist.size()) {
                 hist.push_back(0);
+            }
             hist[pt_level]++;
         }
 
@@ -131,8 +134,9 @@ void hnsw_add_vertices(
             }
 
             // random permutation to get rid of dataset order bias
-            for (int j = i0; j < i1; j++)
+            for (int j = i0; j < i1; j++) {
                 std::swap(order[j], order[j + rng2.rand_int(i1 - j)]);
+            }
 
             bool interrupt = false;
 
@@ -377,8 +381,9 @@ void IndexHNSW::shrink_level_0_neighbors(int new_size) {
 
             for (size_t j = begin; j < end; j++) {
                 int v1 = hnsw.neighbors[j];
-                if (v1 < 0)
+                if (v1 < 0) {
                     break;
+                }
                 initial_list.emplace(dis->symmetric_dis(i, v1), v1);
 
                 // initial_list.emplace(qdis(v1), v1);
@@ -389,10 +394,11 @@ void IndexHNSW::shrink_level_0_neighbors(int new_size) {
                     *dis, initial_list, shrunk_list, new_size);
 
             for (size_t j = begin; j < end; j++) {
-                if (j - begin < shrunk_list.size())
+                if (j - begin < shrunk_list.size()) {
                     hnsw.neighbors[j] = shrunk_list[j - begin].id;
-                else
+                } else {
                     hnsw.neighbors[j] = -1;
+                }
             }
         }
     }
@@ -472,10 +478,12 @@ void IndexHNSW::init_level_0_from_knngraph(
 
         for (size_t j = 0; j < k; j++) {
             int v1 = I[i * k + j];
-            if (v1 == i)
+            if (v1 == i) {
                 continue;
-            if (v1 < 0)
+            }
+            if (v1 < 0) {
                 break;
+            }
             initial_list.emplace(D[i * k + j], v1);
         }
 
@@ -486,10 +494,11 @@ void IndexHNSW::init_level_0_from_knngraph(
         hnsw.neighbor_range(i, 0, &begin, &end);
 
         for (size_t j = begin; j < end; j++) {
-            if (j - begin < shrunk_list.size())
+            if (j - begin < shrunk_list.size()) {
                 hnsw.neighbors[j] = shrunk_list[j - begin].id;
-            else
+            } else {
                 hnsw.neighbors[j] = -1;
+            }
         }
     }
 }
@@ -499,8 +508,9 @@ void IndexHNSW::init_level_0_from_entry_points(
         const storage_idx_t* points,
         const storage_idx_t* nearests) {
     std::vector<omp_lock_t> locks(ntotal);
-    for (int i = 0; i < ntotal; i++)
+    for (int i = 0; i < ntotal; i++) {
         omp_init_lock(&locks[i]);
+    }
 
 #pragma omp parallel
     {
@@ -530,8 +540,9 @@ void IndexHNSW::init_level_0_from_entry_points(
         printf("\n");
     }
 
-    for (int i = 0; i < ntotal; i++)
+    for (int i = 0; i < ntotal; i++) {
         omp_destroy_lock(&locks[i]);
+    }
 }
 
 void IndexHNSW::reorder_links() {
@@ -578,8 +589,9 @@ void IndexHNSW::link_singletons() {
         hnsw.neighbor_range(i, 0, &begin, &end);
         for (size_t j = begin; j < end; j++) {
             storage_idx_t ni = hnsw.neighbors[j];
-            if (ni >= 0)
+            if (ni >= 0) {
                 seen[ni] = true;
+            }
         }
     }
 
@@ -589,8 +601,9 @@ void IndexHNSW::link_singletons() {
         if (!seen[i]) {
             singletons.push_back(i);
             n_sing++;
-            if (hnsw.levels[i] > 1)
+            if (hnsw.levels[i] > 1) {
                 n_sing_l1++;
+            }
         }
     }
 
@@ -722,8 +735,9 @@ int search_from_candidates_2(
 
         for (size_t j = begin; j < end; j++) {
             int v1 = hnsw.neighbors[j];
-            if (v1 < 0)
+            if (v1 < 0) {
                 break;
+            }
             if (vt.visited[v1] == vt.visno + 1) {
                 // nothing to do
             } else {
@@ -749,8 +763,9 @@ int search_from_candidates_2(
     }
 
     stats.n1++;
-    if (candidates.size() == 0)
+    if (candidates.size() == 0) {
         stats.n2++;
+    }
 
     return nres;
 }
@@ -814,8 +829,9 @@ void IndexHNSW2Level::search(
 
                 for (int j = 0; j < nprobe; j++) {
                     idx_t key = coarse_assign[j + i * nprobe];
-                    if (key < 0)
+                    if (key < 0) {
                         break;
+                    }
                     size_t list_length = index_ivfpq->get_list_size(key);
                     const idx_t* ids = index_ivfpq->invlists->get_ids(key);
 
@@ -827,8 +843,9 @@ void IndexHNSW2Level::search(
                 candidates.clear();
 
                 for (int j = 0; j < k; j++) {
-                    if (idxi[j] < 0)
+                    if (idxi[j] < 0) {
                         break;
+                    }
                     candidates.push(idxi[j], simi[j]);
                 }
 
@@ -893,15 +910,31 @@ IndexHNSWCagra::IndexHNSWCagra() {
     is_trained = true;
 }
 
-IndexHNSWCagra::IndexHNSWCagra(int d, int M, MetricType metric)
-        : IndexHNSW(
-                  (metric == METRIC_L2)
-                          ? static_cast<IndexFlat*>(new IndexFlatL2(d))
-                          : static_cast<IndexFlat*>(new IndexFlatIP(d)),
-                  M) {
+IndexHNSWCagra::IndexHNSWCagra(
+        int d,
+        int M,
+        MetricType metric,
+        NumericType numeric_type)
+        : IndexHNSW(d, M, metric) {
     FAISS_THROW_IF_NOT_MSG(
             ((metric == METRIC_L2) || (metric == METRIC_INNER_PRODUCT)),
             "unsupported metric type for IndexHNSWCagra");
+    numeric_type_ = numeric_type;
+    if (numeric_type == NumericType::Float32) {
+        // Use flat storage with full precision for fp32
+        storage = (metric == METRIC_L2)
+                ? static_cast<Index*>(new IndexFlatL2(d))
+                : static_cast<Index*>(new IndexFlatIP(d));
+    } else if (numeric_type == NumericType::Float16) {
+        auto qtype = ScalarQuantizer::QT_fp16;
+        storage = new IndexScalarQuantizer(d, qtype, metric);
+    } else {
+        FAISS_THROW_MSG(
+                "Unsupported numeric_type: only F16 and F32 are supported for IndexHNSWCagra");
+    }
+
+    metric_arg = storage->metric_arg;
+
     own_fields = true;
     is_trained = true;
     init_level0 = true;
@@ -967,4 +1000,12 @@ void IndexHNSWCagra::search(
     }
 }
 
+faiss::NumericType IndexHNSWCagra::get_numeric_type() const {
+    return numeric_type_;
+}
+
+void IndexHNSWCagra::set_numeric_type(faiss::NumericType numeric_type) {
+    numeric_type_ = numeric_type;
+}
+
 } // namespace faiss

data/vendor/faiss/faiss/IndexHNSW.h

@@ -10,6 +10,7 @@
 #pragma once
 
 #include <vector>
+#include "faiss/Index.h"
 
 #include <faiss/IndexFlat.h>
 #include <faiss/IndexPQ.h>
@@ -170,7 +171,11 @@ struct IndexHNSW2Level : IndexHNSW {
 
 struct IndexHNSWCagra : IndexHNSW {
     IndexHNSWCagra();
-    IndexHNSWCagra(int d, int M, MetricType metric = METRIC_L2);
+    IndexHNSWCagra(
+            int d,
+            int M,
+            MetricType metric = METRIC_L2,
+            NumericType numeric_type = NumericType::Float32);
 
     /// When set to true, the index is immutable.
     /// This option is used to copy the knn graph from GpuIndexCagra
@@ -195,6 +200,10 @@ struct IndexHNSWCagra : IndexHNSW {
             float* distances,
             idx_t* labels,
             const SearchParameters* params = nullptr) const override;
+
+    faiss::NumericType get_numeric_type() const;
+    void set_numeric_type(faiss::NumericType numeric_type);
+    NumericType numeric_type_;
 };
 
 } // namespace faiss