faiss 0.2.0 → 0.2.4

data/vendor/faiss/faiss/impl/NSG.cpp

@@ -0,0 +1,679 @@
+/**
+ * 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/impl/NSG.h>
+
+#include <algorithm>
+#include <memory>
+#include <mutex>
+#include <stack>
+
+#include <faiss/impl/AuxIndexStructures.h>
+
+namespace faiss {
+
+namespace nsg {
+
+namespace {
+
+// It needs to be smaller than 0
+constexpr int EMPTY_ID = -1;
+
+/* Wrap the distance computer into one that negates the
+   distances. This makes supporting INNER_PRODUCE search easier */
+
+struct NegativeDistanceComputer : DistanceComputer {
+    using idx_t = Index::idx_t;
+
+    /// 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;
+    }
+};
+
+} // namespace
+
+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 nsg
+
+using namespace nsg;
+
+using LockGuard = std::lock_guard<std::mutex>;
+
+struct Neighbor {
+    int id;
+    float distance;
+    bool flag;
+
+    Neighbor() = default;
+    Neighbor(int id, float distance, bool f)
+            : id(id), distance(distance), flag(f) {}
+
+    inline bool operator<(const Neighbor& other) const {
+        return distance < other.distance;
+    }
+};
+
+struct Node {
+    int id;
+    float distance;
+
+    Node() = default;
+    Node(int id, float distance) : id(id), distance(distance) {}
+
+    inline bool operator<(const Node& other) const {
+        return distance < other.distance;
+    }
+};
+
+inline int insert_into_pool(Neighbor* addr, int K, Neighbor nn) {
+    // find the location to insert
+    int left = 0, right = K - 1;
+    if (addr[left].distance > nn.distance) {
+        memmove(&addr[left + 1], &addr[left], K * sizeof(Neighbor));
+        addr[left] = nn;
+        return left;
+    }
+    if (addr[right].distance < nn.distance) {
+        addr[K] = nn;
+        return K;
+    }
+    while (left < right - 1) {
+        int mid = (left + right) / 2;
+        if (addr[mid].distance > nn.distance) {
+            right = mid;
+        } else {
+            left = mid;
+        }
+    }
+    // check equal ID
+
+    while (left > 0) {
+        if (addr[left].distance < nn.distance) {
+            break;
+        }
+        if (addr[left].id == nn.id) {
+            return K + 1;
+        }
+        left--;
+    }
+    if (addr[left].id == nn.id || addr[right].id == nn.id) {
+        return K + 1;
+    }
+    memmove(&addr[right + 1], &addr[right], (K - right) * sizeof(Neighbor));
+    addr[right] = nn;
+    return right;
+}
+
+NSG::NSG(int R) : R(R), rng(0x0903) {
+    L = R + 32;
+    C = R + 100;
+    search_L = 16;
+    ntotal = 0;
+    is_built = false;
+    srand(0x1998);
+}
+
+void NSG::search(
+        DistanceComputer& dis,
+        int k,
+        idx_t* I,
+        float* D,
+        VisitedTable& vt) const {
+    FAISS_THROW_IF_NOT(is_built);
+    FAISS_THROW_IF_NOT(final_graph);
+
+    int pool_size = std::max(search_L, k);
+    std::vector<Neighbor> retset;
+    std::vector<Node> tmp;
+    search_on_graph<false>(
+            *final_graph, dis, vt, enterpoint, pool_size, retset, tmp);
+
+    for (size_t i = 0; i < k; i++) {
+        I[i] = retset[i].id;
+        D[i] = retset[i].distance;
+    }
+}
+
+void NSG::build(
+        Index* storage,
+        idx_t n,
+        const nsg::Graph<idx_t>& knn_graph,
+        bool verbose) {
+    FAISS_THROW_IF_NOT(!is_built && ntotal == 0);
+
+    if (verbose) {
+        printf("NSG::build R=%d, L=%d, C=%d\n", R, L, C);
+    }
+
+    ntotal = n;
+    init_graph(storage, knn_graph);
+
+    std::vector<int> degrees(n, 0);
+    {
+        nsg::Graph<Node> tmp_graph(n, R);
+
+        link(storage, knn_graph, tmp_graph, verbose);
+
+        final_graph = std::make_shared<nsg::Graph<int>>(n, R);
+        std::fill_n(final_graph->data, n * R, EMPTY_ID);
+
+#pragma omp parallel for
+        for (int i = 0; i < n; i++) {
+            int cnt = 0;
+            for (int j = 0; j < R; j++) {
+                int id = tmp_graph.at(i, j).id;
+                if (id != EMPTY_ID) {
+                    final_graph->at(i, cnt) = id;
+                    cnt += 1;
+                }
+                degrees[i] = cnt;
+            }
+        }
+    }
+
+    int num_attached = tree_grow(storage, degrees);
+    check_graph();
+    is_built = true;
+
+    if (verbose) {
+        int max = 0, min = 1e6;
+        double avg = 0;
+
+        for (int i = 0; i < n; i++) {
+            int size = 0;
+            while (size < R && final_graph->at(i, size) != EMPTY_ID) {
+                size += 1;
+            }
+            max = std::max(size, max);
+            min = std::min(size, min);
+            avg += size;
+        }
+
+        avg = avg / n;
+        printf("Degree Statistics: Max = %d, Min = %d, Avg = %lf\n",
+               max,
+               min,
+               avg);
+        printf("Attached nodes: %d\n", num_attached);
+    }
+}
+
+void NSG::reset() {
+    final_graph.reset();
+    ntotal = 0;
+    is_built = false;
+}
+
+void NSG::init_graph(Index* storage, const nsg::Graph<idx_t>& knn_graph) {
+    int d = storage->d;
+    int n = storage->ntotal;
+
+    std::unique_ptr<float[]> center(new float[d]);
+    std::unique_ptr<float[]> tmp(new float[d]);
+    std::fill_n(center.get(), d, 0.0f);
+
+    for (int i = 0; i < n; i++) {
+        storage->reconstruct(i, tmp.get());
+        for (int j = 0; j < d; j++) {
+            center[j] += tmp[j];
+        }
+    }
+
+    for (int i = 0; i < d; i++) {
+        center[i] /= n;
+    }
+
+    std::vector<Neighbor> retset;
+    std::vector<Node> tmpset;
+
+    // random initialize navigating point
+    int ep = rng.rand_int(n);
+    std::unique_ptr<DistanceComputer> dis(storage_distance_computer(storage));
+
+    dis->set_query(center.get());
+    VisitedTable vt(ntotal);
+
+    // Do not collect the visited nodes
+    search_on_graph<false>(knn_graph, *dis, vt, ep, L, retset, tmpset);
+
+    // set enterpoint
+    enterpoint = retset[0].id;
+}
+
+template <bool collect_fullset, class index_t>
+void NSG::search_on_graph(
+        const nsg::Graph<index_t>& graph,
+        DistanceComputer& dis,
+        VisitedTable& vt,
+        int ep,
+        int pool_size,
+        std::vector<Neighbor>& retset,
+        std::vector<Node>& fullset) const {
+    RandomGenerator gen(0x1234);
+    retset.resize(pool_size + 1);
+    std::vector<int> init_ids(pool_size);
+
+    int num_ids = 0;
+    for (int i = 0; i < init_ids.size() && i < graph.K; i++) {
+        int id = (int)graph.at(ep, i);
+        if (id < 0 || id >= ntotal) {
+            continue;
+        }
+
+        init_ids[i] = id;
+        vt.set(id);
+        num_ids += 1;
+    }
+
+    while (num_ids < pool_size) {
+        int id = gen.rand_int(ntotal);
+        if (vt.get(id)) {
+            continue;
+        }
+
+        init_ids[num_ids] = id;
+        num_ids++;
+        vt.set(id);
+    }
+
+    for (int i = 0; i < init_ids.size(); i++) {
+        int id = init_ids[i];
+
+        float dist = dis(id);
+        retset[i] = Neighbor(id, dist, true);
+
+        if (collect_fullset) {
+            fullset.emplace_back(retset[i].id, retset[i].distance);
+        }
+    }
+
+    std::sort(retset.begin(), retset.begin() + pool_size);
+
+    int k = 0;
+    while (k < pool_size) {
+        int updated_pos = pool_size;
+
+        if (retset[k].flag) {
+            retset[k].flag = false;
+            int n = retset[k].id;
+
+            for (int m = 0; m < graph.K; m++) {
+                int id = (int)graph.at(n, m);
+                if (id < 0 || id > ntotal || vt.get(id)) {
+                    continue;
+                }
+                vt.set(id);
+
+                float dist = dis(id);
+                Neighbor nn(id, dist, true);
+                if (collect_fullset) {
+                    fullset.emplace_back(id, dist);
+                }
+
+                if (dist >= retset[pool_size - 1].distance) {
+                    continue;
+                }
+
+                int r = insert_into_pool(retset.data(), pool_size, nn);
+
+                updated_pos = std::min(updated_pos, r);
+            }
+        }
+
+        k = (updated_pos <= k) ? updated_pos : (k + 1);
+    }
+}
+
+void NSG::link(
+        Index* storage,
+        const nsg::Graph<idx_t>& knn_graph,
+        nsg::Graph<Node>& graph,
+        bool /* verbose */) {
+#pragma omp parallel
+    {
+        std::unique_ptr<float[]> vec(new float[storage->d]);
+
+        std::vector<Node> pool;
+        std::vector<Neighbor> tmp;
+
+        VisitedTable vt(ntotal);
+        std::unique_ptr<DistanceComputer> dis(
+                storage_distance_computer(storage));
+
+#pragma omp for schedule(dynamic, 100)
+        for (int i = 0; i < ntotal; i++) {
+            storage->reconstruct(i, vec.get());
+            dis->set_query(vec.get());
+
+            // Collect the visited nodes into pool
+            search_on_graph<true>(
+                    knn_graph, *dis, vt, enterpoint, L, tmp, pool);
+
+            sync_prune(i, pool, *dis, vt, knn_graph, graph);
+
+            pool.clear();
+            tmp.clear();
+            vt.advance();
+        }
+    } // omp parallel
+
+    std::vector<std::mutex> locks(ntotal);
+#pragma omp parallel
+    {
+        std::unique_ptr<DistanceComputer> dis(
+                storage_distance_computer(storage));
+
+#pragma omp for schedule(dynamic, 100)
+        for (int i = 0; i < ntotal; ++i) {
+            add_reverse_links(i, locks, *dis, graph);
+        }
+    } // omp parallel
+}
+
+void NSG::sync_prune(
+        int q,
+        std::vector<Node>& pool,
+        DistanceComputer& dis,
+        VisitedTable& vt,
+        const nsg::Graph<idx_t>& knn_graph,
+        nsg::Graph<Node>& graph) {
+    for (int i = 0; i < knn_graph.K; i++) {
+        int id = knn_graph.at(q, i);
+        if (id < 0 || id >= ntotal || vt.get(id)) {
+            continue;
+        }
+
+        float dist = dis.symmetric_dis(q, id);
+        pool.emplace_back(id, dist);
+    }
+
+    std::sort(pool.begin(), pool.end());
+
+    std::vector<Node> result;
+
+    int start = 0;
+    if (pool[start].id == q) {
+        start++;
+    }
+    result.push_back(pool[start]);
+
+    while (result.size() < R && (++start) < pool.size() && start < C) {
+        auto& p = pool[start];
+        bool occlude = false;
+        for (int t = 0; t < result.size(); t++) {
+            if (p.id == result[t].id) {
+                occlude = true;
+                break;
+            }
+            float djk = dis.symmetric_dis(result[t].id, p.id);
+            if (djk < p.distance /* dik */) {
+                occlude = true;
+                break;
+            }
+        }
+        if (!occlude) {
+            result.push_back(p);
+        }
+    }
+
+    for (size_t i = 0; i < R; i++) {
+        if (i < result.size()) {
+            graph.at(q, i).id = result[i].id;
+            graph.at(q, i).distance = result[i].distance;
+        } else {
+            graph.at(q, i).id = EMPTY_ID;
+        }
+    }
+}
+
+void NSG::add_reverse_links(
+        int q,
+        std::vector<std::mutex>& locks,
+        DistanceComputer& dis,
+        nsg::Graph<Node>& graph) {
+    for (size_t i = 0; i < R; i++) {
+        if (graph.at(q, i).id == EMPTY_ID) {
+            break;
+        }
+
+        Node sn(q, graph.at(q, i).distance);
+        int des = graph.at(q, i).id;
+
+        std::vector<Node> tmp_pool;
+        int dup = 0;
+        {
+            LockGuard guard(locks[des]);
+            for (int j = 0; j < R; j++) {
+                if (graph.at(des, j).id == EMPTY_ID) {
+                    break;
+                }
+                if (q == graph.at(des, j).id) {
+                    dup = 1;
+                    break;
+                }
+                tmp_pool.push_back(graph.at(des, j));
+            }
+        }
+
+        if (dup) {
+            continue;
+        }
+
+        tmp_pool.push_back(sn);
+        if (tmp_pool.size() > R) {
+            std::vector<Node> result;
+            int start = 0;
+            std::sort(tmp_pool.begin(), tmp_pool.end());
+            result.push_back(tmp_pool[start]);
+
+            while (result.size() < R && (++start) < tmp_pool.size()) {
+                auto& p = tmp_pool[start];
+                bool occlude = false;
+
+                for (int t = 0; t < result.size(); t++) {
+                    if (p.id == result[t].id) {
+                        occlude = true;
+                        break;
+                    }
+                    float djk = dis.symmetric_dis(result[t].id, p.id);
+                    if (djk < p.distance /* dik */) {
+                        occlude = true;
+                        break;
+                    }
+                }
+
+                if (!occlude) {
+                    result.push_back(p);
+                }
+            }
+
+            {
+                LockGuard guard(locks[des]);
+                for (int t = 0; t < result.size(); t++) {
+                    graph.at(des, t) = result[t];
+                }
+            }
+
+        } else {
+            LockGuard guard(locks[des]);
+            for (int t = 0; t < R; t++) {
+                if (graph.at(des, t).id == EMPTY_ID) {
+                    graph.at(des, t) = sn;
+                    break;
+                }
+            }
+        }
+    }
+}
+
+int NSG::tree_grow(Index* storage, std::vector<int>& degrees) {
+    int root = enterpoint;
+    VisitedTable vt(ntotal);
+    VisitedTable vt2(ntotal);
+
+    int num_attached = 0;
+    int cnt = 0;
+    while (true) {
+        cnt = dfs(vt, root, cnt);
+        if (cnt >= ntotal) {
+            break;
+        }
+
+        root = attach_unlinked(storage, vt, vt2, degrees);
+        vt2.advance();
+        num_attached += 1;
+    }
+
+    return num_attached;
+}
+
+int NSG::dfs(VisitedTable& vt, int root, int cnt) const {
+    int node = root;
+    std::stack<int> stack;
+    stack.push(root);
+
+    if (!vt.get(root)) {
+        cnt++;
+    }
+    vt.set(root);
+
+    while (!stack.empty()) {
+        int next = EMPTY_ID;
+        for (int i = 0; i < R; i++) {
+            int id = final_graph->at(node, i);
+            if (id != EMPTY_ID && !vt.get(id)) {
+                next = id;
+                break;
+            }
+        }
+
+        if (next == EMPTY_ID) {
+            stack.pop();
+            if (stack.empty()) {
+                break;
+            }
+            node = stack.top();
+            continue;
+        }
+        node = next;
+        vt.set(node);
+        stack.push(node);
+        cnt++;
+    }
+
+    return cnt;
+}
+
+int NSG::attach_unlinked(
+        Index* storage,
+        VisitedTable& vt,
+        VisitedTable& vt2,
+        std::vector<int>& degrees) {
+    /* NOTE: This implementation is slightly different from the original paper.
+     *
+     * Instead of connecting the unlinked node to the nearest point in the
+     * spanning tree which will increase the maximum degree of the graph and
+     * also make the graph hard to maintain, this implementation links the
+     * unlinked node to the nearest node of which the degree is smaller than R.
+     * It will keep the degree of all nodes to be no more than `R`.
+     */
+
+    // find one unlinked node
+    int id = EMPTY_ID;
+    for (int i = 0; i < ntotal; i++) {
+        if (!vt.get(i)) {
+            id = i;
+            break;
+        }
+    }
+
+    if (id == EMPTY_ID) {
+        return EMPTY_ID; // No Unlinked Node
+    }
+
+    std::vector<Neighbor> tmp;
+    std::vector<Node> pool;
+
+    std::unique_ptr<DistanceComputer> dis(storage_distance_computer(storage));
+    std::unique_ptr<float[]> vec(new float[storage->d]);
+
+    storage->reconstruct(id, vec.get());
+    dis->set_query(vec.get());
+
+    // Collect the visited nodes into pool
+    search_on_graph<true>(
+            *final_graph, *dis, vt2, enterpoint, search_L, tmp, pool);
+
+    std::sort(pool.begin(), pool.end());
+
+    int node;
+    bool found = false;
+    for (int i = 0; i < pool.size(); i++) {
+        node = pool[i].id;
+        if (degrees[node] < R && node != id) {
+            found = true;
+            break;
+        }
+    }
+
+    // randomly choice annother node
+    if (!found) {
+        do {
+            node = rng.rand_int(ntotal);
+            if (vt.get(node) && degrees[node] < R && node != id) {
+                found = true;
+            }
+        } while (!found);
+    }
+
+    int pos = degrees[node];
+    final_graph->at(node, pos) = id; // replace
+    degrees[node] += 1;
+
+    return node;
+}
+
+void NSG::check_graph() const {
+#pragma omp parallel for
+    for (int i = 0; i < ntotal; i++) {
+        for (int j = 0; j < R; j++) {
+            int id = final_graph->at(i, j);
+            FAISS_THROW_IF_NOT(id < ntotal && (id >= 0 || id == EMPTY_ID));
+        }
+    }
+}
+
+} // namespace faiss