faiss 0.1.7 → 0.2.3

data/vendor/faiss/faiss/IndexNNDescent.cpp

@@ -0,0 +1,231 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// -*- c++ -*-
+
+#include <faiss/IndexNNDescent.h>
+
+#include <omp.h>
+
+#include <cinttypes>
+#include <cstdio>
+#include <cstdlib>
+
+#include <queue>
+#include <unordered_set>
+
+#ifdef __SSE__
+#endif
+
+#include <faiss/IndexFlat.h>
+#include <faiss/impl/AuxIndexStructures.h>
+#include <faiss/impl/FaissAssert.h>
+#include <faiss/utils/Heap.h>
+#include <faiss/utils/distances.h>
+#include <faiss/utils/random.h>
+
+extern "C" {
+
+/* declare BLAS functions, see http://www.netlib.org/clapack/cblas/ */
+
+int sgemm_(
+        const char* transa,
+        const char* transb,
+        FINTEGER* m,
+        FINTEGER* n,
+        FINTEGER* k,
+        const float* alpha,
+        const float* a,
+        FINTEGER* lda,
+        const float* b,
+        FINTEGER* ldb,
+        float* beta,
+        float* c,
+        FINTEGER* ldc);
+}
+
+namespace faiss {
+
+using idx_t = Index::idx_t;
+using storage_idx_t = NNDescent::storage_idx_t;
+
+/**************************************************************
+ * add / search blocks of descriptors
+ **************************************************************/
+
+namespace {
+
+/* Wrap the distance computer into one that negates the
+   distances. This makes supporting INNER_PRODUCE search easier */
+
+struct NegativeDistanceComputer : DistanceComputer {
+    /// owned by this
+    DistanceComputer* basedis;
+
+    explicit NegativeDistanceComputer(DistanceComputer* basedis)
+            : basedis(basedis) {}
+
+    void set_query(const float* x) override {
+        basedis->set_query(x);
+    }
+
+    /// compute distance of vector i to current query
+    float operator()(idx_t i) override {
+        return -(*basedis)(i);
+    }
+
+    /// compute distance between two stored vectors
+    float symmetric_dis(idx_t i, idx_t j) override {
+        return -basedis->symmetric_dis(i, j);
+    }
+
+    ~NegativeDistanceComputer() override {
+        delete basedis;
+    }
+};
+
+DistanceComputer* storage_distance_computer(const Index* storage) {
+    if (storage->metric_type == METRIC_INNER_PRODUCT) {
+        return new NegativeDistanceComputer(storage->get_distance_computer());
+    } else {
+        return storage->get_distance_computer();
+    }
+}
+
+} // namespace
+
+/**************************************************************
+ * IndexNNDescent implementation
+ **************************************************************/
+
+IndexNNDescent::IndexNNDescent(int d, int K, MetricType metric)
+        : Index(d, metric),
+          nndescent(d, K),
+          own_fields(false),
+          storage(nullptr) {}
+
+IndexNNDescent::IndexNNDescent(Index* storage, int K)
+        : Index(storage->d, storage->metric_type),
+          nndescent(storage->d, K),
+          own_fields(false),
+          storage(storage) {}
+
+IndexNNDescent::~IndexNNDescent() {
+    if (own_fields) {
+        delete storage;
+    }
+}
+
+void IndexNNDescent::train(idx_t n, const float* x) {
+    FAISS_THROW_IF_NOT_MSG(
+            storage,
+            "Please use IndexNNDescentFlat (or variants) "
+            "instead of IndexNNDescent directly");
+    // nndescent structure does not require training
+    storage->train(n, x);
+    is_trained = true;
+}
+
+void IndexNNDescent::search(
+        idx_t n,
+        const float* x,
+        idx_t k,
+        float* distances,
+        idx_t* labels) const
+
+{
+    FAISS_THROW_IF_NOT_MSG(
+            storage,
+            "Please use IndexNNDescentFlat (or variants) "
+            "instead of IndexNNDescent directly");
+    if (verbose) {
+        printf("Parameters: k=%" PRId64 ", search_L=%d\n",
+               k,
+               nndescent.search_L);
+    }
+
+    idx_t check_period =
+            InterruptCallback::get_period_hint(d * nndescent.search_L);
+
+    for (idx_t i0 = 0; i0 < n; i0 += check_period) {
+        idx_t i1 = std::min(i0 + check_period, n);
+
+#pragma omp parallel
+        {
+            VisitedTable vt(ntotal);
+
+            DistanceComputer* dis = storage_distance_computer(storage);
+            ScopeDeleter1<DistanceComputer> del(dis);
+
+#pragma omp for
+            for (idx_t i = i0; i < i1; i++) {
+                idx_t* idxi = labels + i * k;
+                float* simi = distances + i * k;
+                dis->set_query(x + i * d);
+
+                maxheap_heapify(k, simi, idxi);
+                nndescent.search(*dis, k, idxi, simi, vt);
+                maxheap_reorder(k, simi, idxi);
+            }
+        }
+        InterruptCallback::check();
+    }
+
+    if (metric_type == METRIC_INNER_PRODUCT) {
+        // we need to revert the negated distances
+        for (size_t i = 0; i < k * n; i++) {
+            distances[i] = -distances[i];
+        }
+    }
+}
+
+void IndexNNDescent::add(idx_t n, const float* x) {
+    FAISS_THROW_IF_NOT_MSG(
+            storage,
+            "Please use IndexNNDescentFlat (or variants) "
+            "instead of IndexNNDescent directly");
+    FAISS_THROW_IF_NOT(is_trained);
+
+    if (ntotal != 0) {
+        fprintf(stderr,
+                "WARNING NNDescent doest not support dynamic insertions,"
+                "multiple insertions would lead to re-building the index");
+    }
+
+    storage->add(n, x);
+    ntotal = storage->ntotal;
+
+    DistanceComputer* dis = storage_distance_computer(storage);
+    ScopeDeleter1<DistanceComputer> del(dis);
+    nndescent.build(*dis, ntotal, verbose);
+}
+
+void IndexNNDescent::reset() {
+    nndescent.reset();
+    storage->reset();
+    ntotal = 0;
+}
+
+void IndexNNDescent::reconstruct(idx_t key, float* recons) const {
+    storage->reconstruct(key, recons);
+}
+
+/**************************************************************
+ * IndexNNDescentFlat implementation
+ **************************************************************/
+
+IndexNNDescentFlat::IndexNNDescentFlat() {
+    is_trained = true;
+}
+
+IndexNNDescentFlat::IndexNNDescentFlat(int d, int M, MetricType metric)
+        : IndexNNDescent(new IndexFlat(d, metric), M) {
+    own_fields = true;
+    is_trained = true;
+}
+
+} // namespace faiss

data/vendor/faiss/faiss/IndexNNDescent.h

@@ -0,0 +1,72 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// -*- c++ -*-
+
+#pragma once
+
+#include <vector>
+
+#include <faiss/IndexFlat.h>
+#include <faiss/impl/NNDescent.h>
+#include <faiss/utils/utils.h>
+
+namespace faiss {
+
+/** The NNDescent index is a normal random-access index with an NNDescent
+ * link structure built on top */
+
+struct IndexNNDescent : Index {
+    // internal storage of vectors (32 bits)
+    using storage_idx_t = NNDescent::storage_idx_t;
+
+    /// Faiss results are 64-bit
+    using idx_t = Index::idx_t;
+
+    // the link strcuture
+    NNDescent nndescent;
+
+    // the sequential storage
+    bool own_fields;
+    Index* storage;
+
+    explicit IndexNNDescent(
+            int d = 0,
+            int K = 32,
+            MetricType metric = METRIC_L2);
+    explicit IndexNNDescent(Index* storage, int K = 32);
+
+    ~IndexNNDescent() override;
+
+    void add(idx_t n, const float* x) override;
+
+    /// Trains the storage if needed
+    void train(idx_t n, const float* x) override;
+
+    /// entry point for search
+    void search(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels) const override;
+
+    void reconstruct(idx_t key, float* recons) const override;
+
+    void reset() override;
+};
+
+/** Flat index topped with with a NNDescent structure to access elements
+ *  more efficiently.
+ */
+
+struct IndexNNDescentFlat : IndexNNDescent {
+    IndexNNDescentFlat();
+    IndexNNDescentFlat(int d, int K, MetricType metric = METRIC_L2);
+};
+
+} // namespace faiss

data/vendor/faiss/faiss/IndexNSG.cpp

@@ -0,0 +1,303 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// -*- c++ -*-
+
+#include <faiss/IndexNSG.h>
+
+#include <omp.h>
+
+#include <cinttypes>
+#include <memory>
+
+#include <faiss/IndexFlat.h>
+#include <faiss/IndexNNDescent.h>
+#include <faiss/impl/AuxIndexStructures.h>
+#include <faiss/impl/FaissAssert.h>
+#include <faiss/utils/Heap.h>
+#include <faiss/utils/distances.h>
+
+namespace faiss {
+
+using idx_t = Index::idx_t;
+using namespace nsg;
+
+/**************************************************************
+ * IndexNSG implementation
+ **************************************************************/
+
+IndexNSG::IndexNSG(int d, int R, MetricType metric)
+        : Index(d, metric),
+          nsg(R),
+          own_fields(false),
+          storage(nullptr),
+          is_built(false),
+          GK(64),
+          build_type(0) {
+    nndescent_S = 10;
+    nndescent_R = 100;
+    nndescent_L = GK + 50;
+    nndescent_iter = 10;
+}
+
+IndexNSG::IndexNSG(Index* storage, int R)
+        : Index(storage->d, storage->metric_type),
+          nsg(R),
+          own_fields(false),
+          storage(storage),
+          is_built(false),
+          GK(64),
+          build_type(1) {
+    nndescent_S = 10;
+    nndescent_R = 100;
+    nndescent_L = GK + 50;
+    nndescent_iter = 10;
+}
+
+IndexNSG::~IndexNSG() {
+    if (own_fields) {
+        delete storage;
+    }
+}
+
+void IndexNSG::train(idx_t n, const float* x) {
+    FAISS_THROW_IF_NOT_MSG(
+            storage,
+            "Please use IndexNSGFlat (or variants) instead of IndexNSG directly");
+    // nsg structure does not require training
+    storage->train(n, x);
+    is_trained = true;
+}
+
+void IndexNSG::search(
+        idx_t n,
+        const float* x,
+        idx_t k,
+        float* distances,
+        idx_t* labels) const
+
+{
+    FAISS_THROW_IF_NOT_MSG(
+            storage,
+            "Please use IndexNSGFlat (or variants) instead of IndexNSG directly");
+
+    int L = std::max(nsg.search_L, (int)k); // in case of search L = -1
+    idx_t check_period = InterruptCallback::get_period_hint(d * L);
+
+    for (idx_t i0 = 0; i0 < n; i0 += check_period) {
+        idx_t i1 = std::min(i0 + check_period, n);
+
+#pragma omp parallel
+        {
+            VisitedTable vt(ntotal);
+
+            DistanceComputer* dis = storage_distance_computer(storage);
+            ScopeDeleter1<DistanceComputer> del(dis);
+
+#pragma omp for
+            for (idx_t i = i0; i < i1; i++) {
+                idx_t* idxi = labels + i * k;
+                float* simi = distances + i * k;
+                dis->set_query(x + i * d);
+
+                maxheap_heapify(k, simi, idxi);
+                nsg.search(*dis, k, idxi, simi, vt);
+                maxheap_reorder(k, simi, idxi);
+
+                vt.advance();
+            }
+        }
+        InterruptCallback::check();
+    }
+
+    if (metric_type == METRIC_INNER_PRODUCT) {
+        // we need to revert the negated distances
+        for (size_t i = 0; i < k * n; i++) {
+            distances[i] = -distances[i];
+        }
+    }
+}
+
+void IndexNSG::build(idx_t n, const float* x, idx_t* knn_graph, int GK) {
+    FAISS_THROW_IF_NOT_MSG(
+            storage,
+            "Please use IndexNSGFlat (or variants) instead of IndexNSG directly");
+    FAISS_THROW_IF_NOT_MSG(
+            !is_built && ntotal == 0, "The IndexNSG is already built");
+
+    storage->add(n, x);
+    ntotal = storage->ntotal;
+
+    // check the knn graph
+    check_knn_graph(knn_graph, n, GK);
+
+    const nsg::Graph<idx_t> knng(knn_graph, n, GK);
+    nsg.build(storage, n, knng, verbose);
+    is_built = true;
+}
+
+void IndexNSG::add(idx_t n, const float* x) {
+    FAISS_THROW_IF_NOT_MSG(
+            storage,
+            "Please use IndexNSGFlat (or variants) "
+            "instead of IndexNSG directly");
+    FAISS_THROW_IF_NOT(is_trained);
+
+    FAISS_THROW_IF_NOT_MSG(
+            !is_built && ntotal == 0,
+            "NSG does not support incremental addition");
+
+    std::vector<idx_t> knng;
+    if (verbose) {
+        printf("IndexNSG::add %zd vectors\n", size_t(n));
+    }
+
+    if (build_type == 0) { // build with brute force search
+
+        if (verbose) {
+            printf("  Build knn graph with brute force search on storage index\n");
+        }
+
+        storage->add(n, x);
+        ntotal = storage->ntotal;
+        FAISS_THROW_IF_NOT(ntotal == n);
+
+        knng.resize(ntotal * (GK + 1));
+        storage->assign(ntotal, x, knng.data(), GK + 1);
+
+        // Remove itself
+        // - For metric distance, we just need to remove the first neighbor
+        // - But for non-metric, e.g. inner product, we need to check
+        // - each neighbor
+        if (storage->metric_type == METRIC_INNER_PRODUCT) {
+            for (idx_t i = 0; i < ntotal; i++) {
+                int count = 0;
+                for (int j = 0; j < GK + 1; j++) {
+                    idx_t id = knng[i * (GK + 1) + j];
+                    if (id != i) {
+                        knng[i * GK + count] = id;
+                        count += 1;
+                    }
+                    if (count == GK) {
+                        break;
+                    }
+                }
+            }
+        } else {
+            for (idx_t i = 0; i < ntotal; i++) {
+                memmove(knng.data() + i * GK,
+                        knng.data() + i * (GK + 1) + 1,
+                        GK * sizeof(idx_t));
+            }
+        }
+
+    } else if (build_type == 1) { // build with NNDescent
+        IndexNNDescent index(storage, GK);
+        index.nndescent.S = nndescent_S;
+        index.nndescent.R = nndescent_R;
+        index.nndescent.L = std::max(nndescent_L, GK + 50);
+        index.nndescent.iter = nndescent_iter;
+        index.verbose = verbose;
+
+        if (verbose) {
+            printf("  Build knn graph with NNdescent S=%d R=%d L=%d niter=%d\n",
+                   index.nndescent.S,
+                   index.nndescent.R,
+                   index.nndescent.L,
+                   index.nndescent.iter);
+        }
+
+        // prevent IndexNSG from deleting the storage
+        index.own_fields = false;
+
+        index.add(n, x);
+
+        // storage->add is already implicit called in IndexNSG.add
+        ntotal = storage->ntotal;
+        FAISS_THROW_IF_NOT(ntotal == n);
+
+        knng.resize(ntotal * GK);
+
+        // cast from idx_t to int
+        const int* knn_graph = index.nndescent.final_graph.data();
+#pragma omp parallel for
+        for (idx_t i = 0; i < ntotal * GK; i++) {
+            knng[i] = knn_graph[i];
+        }
+    } else {
+        FAISS_THROW_MSG("build_type should be 0 or 1");
+    }
+
+    if (verbose) {
+        printf("  Check the knn graph\n");
+    }
+
+    // check the knn graph
+    check_knn_graph(knng.data(), n, GK);
+
+    if (verbose) {
+        printf("  nsg building\n");
+    }
+
+    const nsg::Graph<idx_t> knn_graph(knng.data(), n, GK);
+    nsg.build(storage, n, knn_graph, verbose);
+    is_built = true;
+}
+
+void IndexNSG::reset() {
+    nsg.reset();
+    storage->reset();
+    ntotal = 0;
+    is_built = false;
+}
+
+void IndexNSG::reconstruct(idx_t key, float* recons) const {
+    storage->reconstruct(key, recons);
+}
+
+void IndexNSG::check_knn_graph(const idx_t* knn_graph, idx_t n, int K) const {
+    idx_t total_count = 0;
+
+#pragma omp parallel for reduction(+ : total_count)
+    for (idx_t i = 0; i < n; i++) {
+        int count = 0;
+        for (int j = 0; j < K; j++) {
+            idx_t id = knn_graph[i * K + j];
+            if (id < 0 || id >= n || id == i) {
+                count += 1;
+            }
+        }
+        total_count += count;
+    }
+
+    if (total_count > 0) {
+        fprintf(stderr,
+                "WARNING: the input knn graph "
+                "has %" PRId64 " invalid entries\n",
+                total_count);
+    }
+    FAISS_THROW_IF_NOT_MSG(
+            total_count < n / 10,
+            "There are too much invalid entries in the knn graph. "
+            "It may be an invalid knn graph.");
+}
+
+/**************************************************************
+ * IndexNSGFlat implementation
+ **************************************************************/
+
+IndexNSGFlat::IndexNSGFlat() {
+    is_trained = true;
+}
+
+IndexNSGFlat::IndexNSGFlat(int d, int R, MetricType metric)
+        : IndexNSG(new IndexFlat(d, metric), R) {
+    own_fields = true;
+    is_trained = true;
+}
+
+} // namespace faiss

data/vendor/faiss/faiss/IndexNSG.h

@@ -0,0 +1,85 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// -*- c++ -*-
+
+#pragma once
+
+#include <vector>
+
+#include <faiss/IndexFlat.h>
+#include <faiss/IndexNNDescent.h>
+#include <faiss/impl/NSG.h>
+#include <faiss/utils/utils.h>
+
+namespace faiss {
+
+/** The NSG index is a normal random-access index with a NSG
+ * link structure built on top */
+
+struct IndexNSG : Index {
+    /// the link strcuture
+    NSG nsg;
+
+    /// the sequential storage
+    bool own_fields;
+    Index* storage;
+
+    /// the index is built or not
+    bool is_built;
+
+    /// K of KNN graph for building
+    int GK;
+
+    /// indicate how to build a knn graph
+    /// - 0: build NSG with brute force search
+    /// - 1: build NSG with NNDescent
+    char build_type;
+
+    /// parameters for nndescent
+    int nndescent_S;
+    int nndescent_R;
+    int nndescent_L;
+    int nndescent_iter;
+
+    explicit IndexNSG(int d = 0, int R = 32, MetricType metric = METRIC_L2);
+    explicit IndexNSG(Index* storage, int R = 32);
+
+    ~IndexNSG() override;
+
+    void build(idx_t n, const float* x, idx_t* knn_graph, int GK);
+
+    void add(idx_t n, const float* x) override;
+
+    /// Trains the storage if needed
+    void train(idx_t n, const float* x) override;
+
+    /// entry point for search
+    void search(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels) const override;
+
+    void reconstruct(idx_t key, float* recons) const override;
+
+    void reset() override;
+
+    void check_knn_graph(const idx_t* knn_graph, idx_t n, int K) const;
+};
+
+/** Flat index topped with with a NSG structure to access elements
+ *  more efficiently.
+ */
+
+struct IndexNSGFlat : IndexNSG {
+    IndexNSGFlat();
+    IndexNSGFlat(int d, int R, MetricType metric = METRIC_L2);
+};
+
+} // namespace faiss