summaryrefslogtreecommitdiffstats
diff options
context:
space:
mode:
authorGeir Storli <geirst@verizonmedia.com>2020-01-17 16:37:51 +0100
committerGitHub <noreply@github.com>2020-01-17 16:37:51 +0100
commit912caa0b773529dc53015c7de1b4b2a9154c3cd2 (patch)
treed387ed186d3367baeefc6c3d15ece8c3b21be467
parent1b6f9a1a9b2b5fabd933e952f315752b0ae5113f (diff)
parent141995a0f961fad2b1ddc507bf96729d4446550b (diff)
Merge pull request #11817 from vespa-engine/arnej/add-hnsw-algo
Arnej/add hnsw algo
-rw-r--r--eval/src/tests/ann/CMakeLists.txt1
-rw-r--r--eval/src/tests/ann/nns-l2.h37
-rw-r--r--eval/src/tests/ann/nns.h23
-rw-r--r--eval/src/tests/ann/sift_benchmark.cpp88
-rw-r--r--eval/src/tests/ann/xp-annoy-nns.cpp59
-rw-r--r--eval/src/tests/ann/xp-hnswlike-nns.cpp394
-rw-r--r--eval/src/tests/ann/xp-lsh-nns.cpp2
7 files changed, 546 insertions, 58 deletions
diff --git a/eval/src/tests/ann/CMakeLists.txt b/eval/src/tests/ann/CMakeLists.txt
index d82b2311b22..05256d19f00 100644
--- a/eval/src/tests/ann/CMakeLists.txt
+++ b/eval/src/tests/ann/CMakeLists.txt
@@ -4,6 +4,7 @@ vespa_add_executable(eval_sift_benchmark_app
SOURCES
sift_benchmark.cpp
xp-annoy-nns.cpp
+ xp-hnswlike-nns.cpp
xp-lsh-nns.cpp
DEPENDS
vespaeval
diff --git a/eval/src/tests/ann/nns-l2.h b/eval/src/tests/ann/nns-l2.h
index cfa5fed704f..dcad5f1bda6 100644
--- a/eval/src/tests/ann/nns-l2.h
+++ b/eval/src/tests/ann/nns-l2.h
@@ -1,8 +1,36 @@
// Copyright 2019 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
#pragma once
+#include <string.h>
#include <vespa/vespalib/hwaccelrated/iaccelrated.h>
+template <typename T, size_t VLEN>
+static double hw_l2_sq_dist(const T * af, const T * bf, size_t sz)
+{
+ constexpr const size_t OpsPerV = VLEN/sizeof(T);
+ typedef T V __attribute__ ((vector_size (VLEN), aligned(VLEN)));
+
+ const V * a = reinterpret_cast<const V *>(af);
+ const V * b = reinterpret_cast<const V *>(bf);
+
+ V tmp_diff;
+ V tmp_squa;
+ V tmp_sum;
+ memset(&tmp_sum, 0, sizeof(tmp_sum));
+
+ const size_t numOps = sz/OpsPerV;
+ for (size_t i = 0; i < numOps; ++i) {
+ tmp_diff = a[i] - b[i];
+ tmp_squa = tmp_diff * tmp_diff;
+ tmp_sum += tmp_squa;
+ }
+ double sum = 0;
+ for (size_t i = 0; i < OpsPerV; ++i) {
+ sum += tmp_sum[i];
+ }
+ return sum;
+}
+
template <typename FltType = float>
struct L2DistCalc {
vespalib::hwaccelrated::IAccelrated::UP _hw;
@@ -11,7 +39,10 @@ struct L2DistCalc {
using Arr = vespalib::ArrayRef<FltType>;
using ConstArr = vespalib::ConstArrayRef<FltType>;
-
+
+ double product(const FltType *v1, const FltType *v2, size_t sz) {
+ return _hw->dotProduct(v1, v2, sz);
+ }
double product(ConstArr v1, ConstArr v2) {
const FltType *p1 = v1.begin();
const FltType *p2 = v2.begin();
@@ -28,9 +59,7 @@ struct L2DistCalc {
return l2sq(tmp);
}
double l2sq_dist(ConstArr v1, ConstArr v2) {
- std::vector<FltType> tmp;
- tmp.resize(v1.size());
- return l2sq_dist(v1, v2, Arr(tmp));
+ return hw_l2_sq_dist<FltType, 32>(v1.cbegin(), v2.cbegin(), v1.size());
}
};
diff --git a/eval/src/tests/ann/nns.h b/eval/src/tests/ann/nns.h
index 2e6666309bd..79c1aac4379 100644
--- a/eval/src/tests/ann/nns.h
+++ b/eval/src/tests/ann/nns.h
@@ -6,23 +6,28 @@
#include "nns-l2.h"
#include <memory>
+struct SqDist {
+ double distance;
+ explicit SqDist(double d) : distance(d) {}
+};
+
struct NnsHit {
uint32_t docid;
- double sqDistance;
- NnsHit(uint32_t di, double sqD)
- : docid(di), sqDistance(sqD) {}
+ SqDist sq;
+ NnsHit(uint32_t di, SqDist sqD)
+ : docid(di), sq(sqD) {}
};
struct NnsHitComparatorLessDistance {
bool operator() (const NnsHit &lhs, const NnsHit& rhs) const {
- if (lhs.sqDistance > rhs.sqDistance) return false;
- if (lhs.sqDistance < rhs.sqDistance) return true;
+ if (lhs.sq.distance > rhs.sq.distance) return false;
+ if (lhs.sq.distance < rhs.sq.distance) return true;
return (lhs.docid > rhs.docid);
}
};
struct NnsHitComparatorGreaterDistance {
bool operator() (const NnsHit &lhs, const NnsHit& rhs) const {
- if (lhs.sqDistance < rhs.sqDistance) return false;
- if (lhs.sqDistance > rhs.sqDistance) return true;
+ if (lhs.sq.distance < rhs.sq.distance) return false;
+ if (lhs.sq.distance > rhs.sq.distance) return true;
return (lhs.docid > rhs.docid);
}
};
@@ -58,3 +63,7 @@ make_annoy_nns(uint32_t numDims, const DocVectorAccess<float> &dva);
extern
std::unique_ptr<NNS<float>>
make_rplsh_nns(uint32_t numDims, const DocVectorAccess<float> &dva);
+
+extern
+std::unique_ptr<NNS<float>>
+make_hnsw_nns(uint32_t numDims, const DocVectorAccess<float> &dva);
diff --git a/eval/src/tests/ann/sift_benchmark.cpp b/eval/src/tests/ann/sift_benchmark.cpp
index f64351166c1..f20df926f24 100644
--- a/eval/src/tests/ann/sift_benchmark.cpp
+++ b/eval/src/tests/ann/sift_benchmark.cpp
@@ -19,6 +19,7 @@
#include "for-sift-top-k.h"
std::vector<TopK> bruteforceResults;
+std::vector<float> tmp_v(NUM_DIMS);
struct PointVector {
float v[NUM_DIMS];
@@ -26,11 +27,17 @@ struct PointVector {
operator ConstArr() const { return ConstArr(v, NUM_DIMS); }
};
-static PointVector *generatedQueries =
- (PointVector *) malloc(NUM_Q * sizeof(PointVector));
+static PointVector *aligned_alloc(size_t num) {
+ char *mem = (char *)malloc(num * sizeof(PointVector) + 512);
+ mem += 512;
+ size_t val = (size_t)mem;
+ size_t unalign = val % 512;
+ mem -= unalign;
+ return (PointVector *)mem;
+}
-static PointVector *generatedDocs =
- (PointVector *) malloc(NUM_DOCS * sizeof(PointVector));
+static PointVector *generatedQueries = aligned_alloc(NUM_Q);
+static PointVector *generatedDocs = aligned_alloc(NUM_DOCS);
struct DocVectorAdapter : public DocVectorAccess<float>
{
@@ -42,7 +49,7 @@ struct DocVectorAdapter : public DocVectorAccess<float>
double computeDistance(const PointVector &query, uint32_t docid) {
const PointVector &docvector = generatedDocs[docid];
- return l2distCalc.l2sq_dist(query, docvector);
+ return l2distCalc.l2sq_dist(query, docvector, tmp_v);
}
void read_queries(std::string fn) {
@@ -137,7 +144,6 @@ public:
TopK bruteforce_nns(const PointVector &query) {
TopK result;
BfHitHeap heap(result.K);
- std::vector<float> tmp_v(NUM_DIMS);
for (uint32_t docid = 0; docid < NUM_DOCS; ++docid) {
const PointVector &docvector = generatedDocs[docid];
double d = l2distCalc.l2sq_dist(query, docvector, tmp_v);
@@ -173,6 +179,7 @@ void verifyBF(uint32_t qid) {
TEST("require that brute force works") {
TimePoint bef = std::chrono::steady_clock::now();
+ bruteforceResults.reserve(NUM_Q);
for (uint32_t cnt = 0; cnt < NUM_Q; ++cnt) {
const PointVector &query = generatedQueries[cnt];
bruteforceResults.emplace_back(bruteforce_nns(query));
@@ -193,7 +200,7 @@ TopK find_with_nns(uint32_t sk, NNS_API &nns, uint32_t qid) {
vespalib::ConstArrayRef<float> query(qv.v, NUM_DIMS);
auto rv = nns.topK(result.K, query, sk);
for (size_t i = 0; i < result.K; ++i) {
- result.hits[i] = Hit(rv[i].docid, rv[i].sqDistance);
+ result.hits[i] = Hit(rv[i].docid, rv[i].sq.distance);
}
return result;
}
@@ -207,6 +214,10 @@ void verify_nns_quality(uint32_t sk, NNS_API &nns, uint32_t qid) {
double c_factor = 1.0;
for (size_t i = 0; i < result.K; ++i) {
double factor = (result.hits[i].distance / perfect.hits[i].distance);
+ if (factor < 0.99 || factor > 25) {
+ fprintf(stderr, "hit[%zu] got distance %.3f, expected %.3f\n",
+ i, result.hits[i].distance, perfect.hits[i].distance);
+ }
sum_error += factor;
c_factor = std::max(c_factor, factor);
}
@@ -223,57 +234,56 @@ void verify_nns_quality(uint32_t sk, NNS_API &nns, uint32_t qid) {
}
}
-TEST("require that Locality Sensitive Hashing mostly works") {
+void benchmark_nns(const char *name, NNS_API &nns, std::vector<uint32_t> sk_list) {
+ fprintf(stderr, "trying %s indexing...\n", name);
TimePoint bef = std::chrono::steady_clock::now();
- DocVectorAdapter adapter;
- std::unique_ptr<NNS_API> nns = make_rplsh_nns(NUM_DIMS, adapter);
for (uint32_t i = 0; i < NUM_DOCS; ++i) {
- nns->addDoc(i);
+ nns.addDoc(i);
}
fprintf(stderr, "added %u documents...\n", NUM_DOCS);
+ find_with_nns(1, nns, 0);
TimePoint aft = std::chrono::steady_clock::now();
- fprintf(stderr, "build RPLSH index: %.3f ms\n", to_ms(aft - bef));
+ fprintf(stderr, "build %s index: %.3f ms\n", name, to_ms(aft - bef));
- for (uint32_t search_k : { 200, 1000 }) {
+ for (uint32_t search_k : sk_list) {
bef = std::chrono::steady_clock::now();
for (int cnt = 0; cnt < NUM_Q; ++cnt) {
- find_with_nns(search_k, *nns, cnt);
+ find_with_nns(search_k, nns, cnt);
}
aft = std::chrono::steady_clock::now();
- fprintf(stderr, "timing for RPLSH search_k=%u: %.3f ms = %.3f ms per query\n",
- search_k, to_ms(aft - bef), to_ms(aft - bef)/NUM_Q);
+ fprintf(stderr, "timing for %s search_k=%u: %.3f ms = %.3f ms/q\n",
+ name, search_k, to_ms(aft - bef), to_ms(aft - bef)/NUM_Q);
for (int cnt = 0; cnt < NUM_Q; ++cnt) {
- verify_nns_quality(search_k, *nns, cnt);
+ verify_nns_quality(search_k, nns, cnt);
}
}
}
-TEST("require that Indexing via NNS api mostly works") {
- fprintf(stderr, "trying indexing...\n");
- TimePoint bef = std::chrono::steady_clock::now();
+
+#if 1
+TEST("require that Locality Sensitive Hashing mostly works") {
+ DocVectorAdapter adapter;
+ std::unique_ptr<NNS_API> nns = make_rplsh_nns(NUM_DIMS, adapter);
+ benchmark_nns("RPLSH", *nns, { 200, 1000 });
+}
+#endif
+
+#if 1
+TEST("require that Annoy via NNS api mostly works") {
DocVectorAdapter adapter;
std::unique_ptr<NNS_API> nns = make_annoy_nns(NUM_DIMS, adapter);
- for (uint32_t i = 0; i < NUM_DOCS; ++i) {
- nns->addDoc(i);
- }
- fprintf(stderr, "added %u documents...\n", NUM_DOCS);
- nns->topK(1, adapter.get(0), 1);
- TimePoint aft = std::chrono::steady_clock::now();
- fprintf(stderr, "build annoy index: %.3f ms\n", to_ms(aft - bef));
+ benchmark_nns("Annoy", *nns, { 8000, 10000 });
+}
+#endif
- for (uint32_t search_k : { 10000, 20000 }) {
- bef = std::chrono::steady_clock::now();
- for (int cnt = 0; cnt < NUM_Q; ++cnt) {
- find_with_nns(search_k, *nns, cnt);
- }
- aft = std::chrono::steady_clock::now();
- fprintf(stderr, "timing for index search_k=%u: %.3f ms = %.3f ms per query\n",
- search_k, to_ms(aft - bef), to_ms(aft - bef)/NUM_Q);
- for (int cnt = 0; cnt < NUM_Q; ++cnt) {
- verify_nns_quality(search_k, *nns, cnt);
- }
- }
+#if 1
+TEST("require that HNSW via NNS api mostly works") {
+ DocVectorAdapter adapter;
+ std::unique_ptr<NNS_API> nns = make_hnsw_nns(NUM_DIMS, adapter);
+ benchmark_nns("HNSW", *nns, { 100, 200 });
}
+#endif
+
int main(int argc, char **argv) {
TEST_MASTER.init(__FILE__);
diff --git a/eval/src/tests/ann/xp-annoy-nns.cpp b/eval/src/tests/ann/xp-annoy-nns.cpp
index e5661c0c044..45392084c80 100644
--- a/eval/src/tests/ann/xp-annoy-nns.cpp
+++ b/eval/src/tests/ann/xp-annoy-nns.cpp
@@ -11,6 +11,12 @@ using V = vespalib::ConstArrayRef<float>;
class AnnoyLikeNns;
struct Node;
+static uint64_t plane_dist_cnt = 0;
+static uint64_t w_cen_dist_cnt = 0;
+static uint64_t leaf_split_cnt = 0;
+static uint64_t find_top_k_cnt = 0;
+static uint64_t find_cand_cnt = 0;
+
using QueueNode = std::pair<double, Node *>;
using NodeQueue = std::priority_queue<QueueNode>;
@@ -20,6 +26,7 @@ struct Node {
virtual Node *addDoc(uint32_t docid, V vector, AnnoyLikeNns &meta) = 0;
virtual int remove(uint32_t docid, V vector) = 0;
virtual void findCandidates(std::set<uint32_t> &cands, V vector, NodeQueue &queue, double minDist) const = 0;
+ virtual void stats(std::vector<uint32_t> &depths) = 0;
};
struct LeafNode : public Node {
@@ -32,6 +39,7 @@ struct LeafNode : public Node {
void findCandidates(std::set<uint32_t> &cands, V vector, NodeQueue &queue, double minDist) const override;
Node *split(AnnoyLikeNns &meta);
+ virtual void stats(std::vector<uint32_t> &depths) override { depths.push_back(1); }
};
struct SplitNode : public Node {
@@ -48,6 +56,12 @@ struct SplitNode : public Node {
void findCandidates(std::set<uint32_t> &cands, V vector, NodeQueue &queue, double minDist) const override;
double planeDistance(V vector) const;
+ virtual void stats(std::vector<uint32_t> &depths) override {
+ size_t i = depths.size();
+ leftChildren->stats(depths);
+ rightChildren->stats(depths);
+ while (i < depths.size()) { ++depths[i++]; }
+ }
};
class AnnoyLikeNns : public NNS<float>
@@ -67,7 +81,10 @@ public:
}
}
+ void dumpStats();
+
~AnnoyLikeNns() {
+ dumpStats();
for (Node *root : _roots) {
delete root;
}
@@ -97,11 +114,9 @@ public:
double
SplitNode::planeDistance(V vector) const
{
+ ++plane_dist_cnt;
assert(vector.size() == hyperPlane.size());
- double dp = 0.0;
- for (size_t i = 0; i < vector.size(); ++i) {
- dp += vector[i] * hyperPlane[i];
- }
+ double dp = l2distCalc.product(&vector[0], &hyperPlane[0], vector.size());
return dp - offsetFromOrigo;
}
@@ -167,6 +182,7 @@ struct WeightedCentroid {
return r;
}
double weightedDistance(V vector) {
+ ++w_cen_dist_cnt;
size_t sz = vector.size();
for (size_t i = 0; i < sz; ++i) {
tmp_vector[i] = vector[i] * cnt;
@@ -179,6 +195,7 @@ struct WeightedCentroid {
Node *
LeafNode::split(AnnoyLikeNns &meta)
{
+ ++leaf_split_cnt;
uint32_t dims = meta.dims();
uint32_t retries = 3;
retry:
@@ -232,6 +249,8 @@ retry:
std::vector<uint32_t> leftDs;
std::vector<uint32_t> rightDs;
+ leftDs.reserve(128);
+ rightDs.reserve(128);
for (uint32_t docid : docids) {
V vector = meta.getVector(docid);
@@ -260,9 +279,9 @@ retry:
#endif
LeafNode *newRightNode = new LeafNode();
- newRightNode->docids = rightDs;
+ newRightNode->docids = std::move(rightDs);
s->rightChildren = newRightNode;
- this->docids = leftDs;
+ this->docids = std::move(leftDs);
s->leftChildren = this;
return s;
}
@@ -327,6 +346,7 @@ SplitNode::findCandidates(std::set<uint32_t> &, V vector, NodeQueue &queue, doub
std::vector<NnsHit>
AnnoyLikeNns::topK(uint32_t k, Vector vector, uint32_t search_k)
{
+ ++find_top_k_cnt;
std::vector<float> tmp;
tmp.resize(_numDims);
vespalib::ArrayRef<float> tmpArr(tmp);
@@ -347,6 +367,7 @@ AnnoyLikeNns::topK(uint32_t k, Vector vector, uint32_t search_k)
Node *n = top.second;
queue.pop();
n->findCandidates(candidates, vector, queue, md);
+ ++find_cand_cnt;
}
#if 0
while (queue.size() > 0) {
@@ -357,7 +378,7 @@ AnnoyLikeNns::topK(uint32_t k, Vector vector, uint32_t search_k)
#endif
for (uint32_t docid : candidates) {
double dist = l2distCalc.l2sq_dist(vector, _dva.get(docid), tmpArr);
- NnsHit hit(docid, dist);
+ NnsHit hit(docid, SqDist(dist));
r.push_back(hit);
}
std::sort(r.begin(), r.end(), NnsHitComparatorLessDistance());
@@ -365,6 +386,30 @@ AnnoyLikeNns::topK(uint32_t k, Vector vector, uint32_t search_k)
return r;
}
+void
+AnnoyLikeNns::dumpStats() {
+ fprintf(stderr, "stats for AnnoyLikeNns:\n");
+ fprintf(stderr, "planeDistance() calls: %zu\n", plane_dist_cnt);
+ fprintf(stderr, "weightedDistance() calls: %zu\n", w_cen_dist_cnt);
+ fprintf(stderr, "leaf split() calls: %zu\n", leaf_split_cnt);
+ fprintf(stderr, "topK() calls: %zu\n", find_top_k_cnt);
+ fprintf(stderr, "findCandidates() calls: %zu\n", find_cand_cnt);
+ std::vector<uint32_t> depths;
+ _roots[0]->stats(depths);
+ std::vector<uint32_t> counts;
+ for (uint32_t deep : depths) {
+ while (counts.size() <= deep) counts.push_back(0);
+ counts[deep]++;
+ }
+ fprintf(stderr, "depths for %zu leaves [\n", depths.size());
+ for (uint32_t deep = 0; deep < counts.size(); ++deep) {
+ if (counts[deep] > 0) {
+ fprintf(stderr, "%u deep count %u\n", deep, counts[deep]);
+ }
+ }
+ fprintf(stderr, "]\n");
+}
+
std::unique_ptr<NNS<float>>
make_annoy_nns(uint32_t numDims, const DocVectorAccess<float> &dva)
{
diff --git a/eval/src/tests/ann/xp-hnswlike-nns.cpp b/eval/src/tests/ann/xp-hnswlike-nns.cpp
new file mode 100644
index 00000000000..ec831610a71
--- /dev/null
+++ b/eval/src/tests/ann/xp-hnswlike-nns.cpp
@@ -0,0 +1,394 @@
+// Copyright 2020 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include <algorithm>
+#include <assert.h>
+#include <queue>
+#include <random>
+#include "nns.h"
+
+using LinkList = std::vector<uint32_t>;
+
+struct Node {
+ std::vector<LinkList> _links;
+ Node(uint32_t , uint32_t numLevels, uint32_t M)
+ : _links(numLevels)
+ {
+ for (uint32_t i = 0; i < _links.size(); ++i) {
+ _links[i].reserve((i == 0) ? (2 * M + 1) : (M+1));
+ }
+ }
+};
+
+struct VisitedSet
+{
+ using Mark = unsigned short;
+ Mark *ptr;
+ Mark curval;
+ size_t sz;
+ VisitedSet(const VisitedSet &) = delete;
+ VisitedSet& operator=(const VisitedSet &) = delete;
+ explicit VisitedSet(size_t size) {
+ ptr = (Mark *)malloc(size * sizeof(Mark));
+ curval = -1;
+ sz = size;
+ }
+ void clear() {
+ ++curval;
+ if (curval == 0) {
+ memset(ptr, 0, sz * sizeof(Mark));
+ ++curval;
+ }
+ }
+ ~VisitedSet() { free(ptr); }
+ void mark(size_t id) { ptr[id] = curval; }
+ bool isMarked(size_t id) const { return ptr[id] == curval; }
+};
+
+struct VisitedSetPool
+{
+ std::unique_ptr<VisitedSet> lastUsed;
+ VisitedSetPool() {
+ lastUsed = std::make_unique<VisitedSet>(250);
+ }
+ ~VisitedSetPool() {}
+ VisitedSet &get(size_t size) {
+ if (size > lastUsed->sz) {
+ lastUsed = std::make_unique<VisitedSet>(size*2);
+ }
+ lastUsed->clear();
+ return *lastUsed;
+ }
+};
+
+struct HnswHit {
+ float dist;
+ uint32_t docid;
+ HnswHit(uint32_t di, SqDist sq) : dist(sq.distance), docid(di) {}
+};
+
+
+using QueueEntry = HnswHit;
+struct GreaterDist {
+ bool operator() (const QueueEntry &lhs, const QueueEntry& rhs) const {
+ return (rhs.dist < lhs.dist);
+ }
+};
+struct LesserDist {
+ bool operator() (const QueueEntry &lhs, const QueueEntry& rhs) const {
+ return (lhs.dist < rhs.dist);
+ }
+};
+
+using NearestList = std::vector<QueueEntry>;
+
+struct NearestPriQ : std::priority_queue<QueueEntry, NearestList, GreaterDist>
+{
+};
+
+struct FurthestPriQ : std::priority_queue<QueueEntry, NearestList, LesserDist>
+{
+ NearestList steal() {
+ NearestList result;
+ c.swap(result);
+ return result;
+ }
+ const NearestList& peek() const { return c; }
+};
+
+class HnswLikeNns : public NNS<float>
+{
+private:
+ std::vector<Node> _nodes;
+ uint32_t _entryId;
+ int _entryLevel;
+ uint32_t _M;
+ uint32_t _efConstruction;
+ double _levelMultiplier;
+ std::default_random_engine _rndGen;
+ VisitedSetPool _visitedSetPool;
+
+ double distance(Vector v, uint32_t id) const;
+
+ double distance(uint32_t a, uint32_t b) const {
+ Vector v = _dva.get(a);
+ return distance(v, b);
+ }
+
+ int randomLevel() {
+ std::uniform_real_distribution<double> distribution(0.0, 1.0);
+ double r = -log(distribution(_rndGen)) * _levelMultiplier;
+ return (int) r;
+ }
+
+public:
+ HnswLikeNns(uint32_t numDims, const DocVectorAccess<float> &dva)
+ : NNS(numDims, dva),
+ _nodes(),
+ _entryId(0),
+ _entryLevel(-1),
+ _M(16),
+ _efConstruction(150),
+ _levelMultiplier(1.0 / log(1.0 * _M))
+ {
+ _nodes.reserve(1234567);
+ }
+
+ ~HnswLikeNns() {}
+
+ LinkList& getLinkList(uint32_t docid, uint32_t level) {
+ // assert(docid < _nodes.size());
+ // assert(level < _nodes[docid]._links.size());
+ return _nodes[docid]._links[level];
+ }
+
+ // simple greedy search
+ QueueEntry search_layer_simple(Vector vector, QueueEntry curPoint, uint32_t searchLevel) {
+ bool keepGoing = true;
+ while (keepGoing) {
+ keepGoing = false;
+ const LinkList& neighbors = getLinkList(curPoint.docid, searchLevel);
+ for (uint32_t n_id : neighbors) {
+ double dist = distance(vector, n_id);
+ if (dist < curPoint.dist) {
+ curPoint = QueueEntry(n_id, SqDist(dist));
+ keepGoing = true;
+ }
+ }
+ }
+ return curPoint;
+ }
+
+ void search_layer_foradd(Vector vector, FurthestPriQ &w,
+ uint32_t ef, uint32_t searchLevel);
+
+ FurthestPriQ search_layer(Vector vector, NearestList entryPoints,
+ uint32_t ef, uint32_t searchLevel) {
+ VisitedSet &visited = _visitedSetPool.get(_nodes.size());
+ NearestPriQ candidates;
+ FurthestPriQ w;
+ for (auto point : entryPoints) {
+ candidates.push(point);
+ w.push(point);
+ visited.mark(point.docid);
+ }
+ double limd = std::numeric_limits<double>::max();
+ while (! candidates.empty()) {
+ QueueEntry cand = candidates.top();
+ candidates.pop();
+ if (cand.dist > limd) {
+ break;
+ }
+ for (uint32_t e_id : getLinkList(cand.docid, searchLevel)) {
+ if (visited.isMarked(e_id)) continue;
+ visited.mark(e_id);
+ double e_dist = distance(vector, e_id);
+ if (e_dist < limd) {
+ candidates.emplace(e_id, SqDist(e_dist));
+ w.emplace(e_id, SqDist(e_dist));
+ if (w.size() > ef) {
+ w.pop();
+ limd = w.top().dist;
+ }
+ }
+ }
+ }
+ return w;
+ }
+
+ bool haveCloserDistance(QueueEntry e, const LinkList &r) const {
+ for (uint32_t prevId : r) {
+ double dist = distance(e.docid, prevId);
+ if (dist < e.dist) return true;
+ }
+ return false;
+ }
+
+ LinkList select_neighbors(NearestPriQ &&w, uint32_t curMax) const;
+
+ LinkList select_neighbors(const NearestList &neighbors, uint32_t curMax) {
+ if (neighbors.size() <= curMax) {
+ LinkList result;
+ result.reserve(curMax+1);
+ for (const auto & entry : neighbors) {
+ result.push_back(entry.docid);
+ }
+ return result;
+ }
+ NearestPriQ w;
+ for (const QueueEntry & entry : neighbors) {
+ w.push(entry);
+ }
+ return select_neighbors(std::move(w), curMax);
+ }
+
+ void addDoc(uint32_t docid) override {
+ Vector vector = _dva.get(docid);
+ for (uint32_t id = _nodes.size(); id <= docid; ++id) {
+ _nodes.emplace_back(id, 0, _M);
+ }
+ int level = randomLevel();
+ assert(_nodes[docid]._links.size() == 0);
+ _nodes[docid] = Node(docid, level+1, _M);
+ if (_entryLevel < 0) {
+ _entryId = docid;
+ _entryLevel = level;
+ return;
+ }
+ int searchLevel = _entryLevel;
+ double entryDist = distance(vector, _entryId);
+ QueueEntry entryPoint(_entryId, SqDist(entryDist));
+ while (searchLevel > level) {
+ entryPoint = search_layer_simple(vector, entryPoint, searchLevel);
+ --searchLevel;
+ }
+ searchLevel = std::min(level, _entryLevel);
+ FurthestPriQ w;
+ w.push(entryPoint);
+ while (searchLevel >= 0) {
+ search_layer_foradd(vector, w, _efConstruction, searchLevel);
+ uint32_t maxLinks = (searchLevel > 0) ? _M : (2 * _M);
+ LinkList neighbors = select_neighbors(w.peek(), maxLinks);
+ connect_new_node(docid, neighbors, searchLevel);
+ each_shrink_ifneeded(neighbors, searchLevel);
+ --searchLevel;
+ }
+ if (level > _entryLevel) {
+ _entryLevel = level;
+ _entryId = docid;
+ }
+ }
+
+ void connect_new_node(uint32_t id, const LinkList &neighbors, uint32_t level);
+
+ void each_shrink_ifneeded(const LinkList &neighbors, uint32_t level);
+
+ void removeDoc(uint32_t ) override {
+ }
+
+ std::vector<NnsHit> topK(uint32_t k, Vector vector, uint32_t search_k) override {
+ std::vector<NnsHit> result;
+ if (_entryLevel < 0) return result;
+ double entryDist = distance(vector, _entryId);
+ QueueEntry entryPoint(_entryId, SqDist(entryDist));
+ int searchLevel = _entryLevel;
+ while (searchLevel > 0) {
+ entryPoint = search_layer_simple(vector, entryPoint, searchLevel);
+ --searchLevel;
+ }
+ NearestList entryPoints;
+ entryPoints.push_back(entryPoint);
+ FurthestPriQ w = search_layer(vector, entryPoints, std::max(k, search_k), 0);
+ if (w.size() < k) {
+ fprintf(stderr, "fewer than expected hits: k=%u, ks=%u, got=%zu\n",
+ k, search_k, w.size());
+ }
+ while (w.size() > k) {
+ w.pop();
+ }
+ std::vector<QueueEntry> reversed;
+ reversed.reserve(w.size());
+ while (! w.empty()) {
+ reversed.push_back(w.top());
+ w.pop();
+ }
+ result.reserve(reversed.size());
+ while (! reversed.empty()) {
+ const QueueEntry &hit = reversed.back();
+ result.emplace_back(hit.docid, SqDist(hit.dist));
+ reversed.pop_back();
+ }
+ return result;
+ }
+};
+
+double
+HnswLikeNns::distance(Vector v, uint32_t b) const
+{
+ Vector w = _dva.get(b);
+ return l2distCalc.l2sq_dist(v, w);
+}
+
+void
+HnswLikeNns::each_shrink_ifneeded(const LinkList &neighbors, uint32_t level) {
+ uint32_t maxLinks = (level > 0) ? _M : (2 * _M);
+ for (uint32_t old_id : neighbors) {
+ LinkList &oldLinks = getLinkList(old_id, level);
+ if (oldLinks.size() > maxLinks) {
+ NearestPriQ w;
+ for (uint32_t n_id : oldLinks) {
+ double n_dist = distance(old_id, n_id);
+ w.emplace(n_id, SqDist(n_dist));
+ }
+ oldLinks = select_neighbors(std::move(w), maxLinks);
+ }
+ }
+}
+
+void
+HnswLikeNns::search_layer_foradd(Vector vector, FurthestPriQ &w,
+ uint32_t ef, uint32_t searchLevel)
+{
+ NearestPriQ candidates;
+ VisitedSet &visited = _visitedSetPool.get(_nodes.size());
+
+ for (const QueueEntry& entry : w.peek()) {
+ candidates.push(entry);
+ visited.mark(entry.docid);
+ }
+
+ double limd = std::numeric_limits<double>::max();
+ while (! candidates.empty()) {
+ QueueEntry cand = candidates.top();
+ candidates.pop();
+ if (cand.dist > limd) {
+ break;
+ }
+ for (uint32_t e_id : getLinkList(cand.docid, searchLevel)) {
+ if (visited.isMarked(e_id)) continue;
+ visited.mark(e_id);
+ double e_dist = distance(vector, e_id);
+ if (e_dist < limd) {
+ candidates.emplace(e_id, SqDist(e_dist));
+ w.emplace(e_id, SqDist(e_dist));
+ if (w.size() > ef) {
+ w.pop();
+ limd = w.top().dist;
+ }
+ }
+ }
+ }
+ return;
+}
+
+LinkList
+HnswLikeNns::select_neighbors(NearestPriQ &&w, uint32_t curMax) const {
+ LinkList result;
+ result.reserve(curMax+1);
+ while (! w.empty()) {
+ QueueEntry e = w.top();
+ w.pop();
+ if (haveCloserDistance(e, result)) continue;
+ result.push_back(e.docid);
+ if (result.size() >= curMax) break;
+ }
+ return result;
+}
+
+void
+HnswLikeNns::connect_new_node(uint32_t id, const LinkList &neighbors, uint32_t level) {
+ LinkList &newLinks = getLinkList(id, level);
+ for (uint32_t neigh_id : neighbors) {
+ LinkList &oldLinks = getLinkList(neigh_id, level);
+ newLinks.push_back(neigh_id);
+ oldLinks.push_back(id);
+ }
+}
+
+
+std::unique_ptr<NNS<float>>
+make_hnsw_nns(uint32_t numDims, const DocVectorAccess<float> &dva)
+{
+ return std::make_unique<HnswLikeNns>(numDims, dva);
+}
+
+
diff --git a/eval/src/tests/ann/xp-lsh-nns.cpp b/eval/src/tests/ann/xp-lsh-nns.cpp
index 285985167c0..0ea119a9c70 100644
--- a/eval/src/tests/ann/xp-lsh-nns.cpp
+++ b/eval/src/tests/ann/xp-lsh-nns.cpp
@@ -231,7 +231,7 @@ RpLshNns::topK(uint32_t k, Vector vector, uint32_t search_k)
std::vector<LshHit> best = heap.bestLshHits();
size_t numHits = std::min((size_t)k, best.size());
for (size_t i = 0; i < numHits; ++i) {
- result.emplace_back(best[i].docid, best[i].distance);
+ result.emplace_back(best[i].docid, SqDist(best[i].distance));
}
return result;
}