diff options
author | Geir Storli <geirst@verizonmedia.com> | 2020-01-17 16:37:51 +0100 |
---|---|---|
committer | GitHub <noreply@github.com> | 2020-01-17 16:37:51 +0100 |
commit | 912caa0b773529dc53015c7de1b4b2a9154c3cd2 (patch) | |
tree | d387ed186d3367baeefc6c3d15ece8c3b21be467 | |
parent | 1b6f9a1a9b2b5fabd933e952f315752b0ae5113f (diff) | |
parent | 141995a0f961fad2b1ddc507bf96729d4446550b (diff) |
Merge pull request #11817 from vespa-engine/arnej/add-hnsw-algo
Arnej/add hnsw algo
-rw-r--r-- | eval/src/tests/ann/CMakeLists.txt | 1 | ||||
-rw-r--r-- | eval/src/tests/ann/nns-l2.h | 37 | ||||
-rw-r--r-- | eval/src/tests/ann/nns.h | 23 | ||||
-rw-r--r-- | eval/src/tests/ann/sift_benchmark.cpp | 88 | ||||
-rw-r--r-- | eval/src/tests/ann/xp-annoy-nns.cpp | 59 | ||||
-rw-r--r-- | eval/src/tests/ann/xp-hnswlike-nns.cpp | 394 | ||||
-rw-r--r-- | eval/src/tests/ann/xp-lsh-nns.cpp | 2 |
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; } |