diff --git a/include/rabitqlib/index/hnsw/hnsw.hpp b/include/rabitqlib/index/hnsw/hnsw.hpp index 3f53b8d..7e290cc 100644 --- a/include/rabitqlib/index/hnsw/hnsw.hpp +++ b/include/rabitqlib/index/hnsw/hnsw.hpp @@ -777,7 +777,7 @@ inline void HierarchicalNSW::add_point( inline maxheap> HierarchicalNSW::search_base_layer( PID ep_id, PID cur_c, int layer ) { - HashBasedBooleanSet* vl = visited_list_pool_->get_free_vislist(); + VisitedSet* vl = visited_list_pool_->get_free_vislist(); maxheap> top_candidates; minheap> candidate_set; @@ -1225,7 +1225,7 @@ inline void HierarchicalNSW::searchBaseLayerST_AdaptiveRerankOptDirect( [[maybe_unused]] const float* query, BoundedKNN& boundedKNN ) { - HashBasedBooleanSet* vl = visited_list_pool_->get_free_vislist(); + VisitedSet* vl = visited_list_pool_->get_free_vislist(); // Use our bounded priority queue instead of the maxheap. buffer::SearchBuffer candidate_set(ef); diff --git a/include/rabitqlib/index/symqg/qg.hpp b/include/rabitqlib/index/symqg/qg.hpp index 3ad41dc..9f91910 100644 --- a/include/rabitqlib/index/symqg/qg.hpp +++ b/include/rabitqlib/index/symqg/qg.hpp @@ -20,12 +20,12 @@ #include "rabitqlib/quantization/rabitq.hpp" #include "rabitqlib/utils/array.hpp" #include "rabitqlib/utils/buffer.hpp" -#include "rabitqlib/utils/hashset.hpp" #include "rabitqlib/utils/io.hpp" #include "rabitqlib/utils/memory.hpp" #include "rabitqlib/utils/rotator.hpp" #include "rabitqlib/utils/space.hpp" #include "rabitqlib/utils/visited_pool.hpp" +#include "rabitqlib/utils/visited_set.hpp" namespace rabitqlib::symqg { @@ -94,25 +94,16 @@ class QuantizedGraph { ); } - void find_candidates( - PID, - size_t, - std::vector>&, - HashBasedBooleanSet&, - const std::vector& - ) const; + void + find_candidates(PID, size_t, std::vector>&, VisitedSet&, const std::vector&) + const; void update_qg(PID, const std::vector>&); - void update_results(buffer::SearchBuffer&, HashBasedBooleanSet&, const T*); + void update_results(buffer::SearchBuffer&, VisitedSet&, const T*); void scan_neighbors( - const BatchQuery&, - PID, - T*, - buffer::SearchBuffer&, - HashBasedBooleanSet&, - size_t + const BatchQuery&, PID, T*, buffer::SearchBuffer&, VisitedSet&, size_t ) const; public: @@ -352,7 +343,7 @@ void QuantizedGraph::scan_neighbors( PID data_id, T* est_dist, buffer::SearchBuffer& search_pool, - HashBasedBooleanSet& vis, + VisitedSet& vis, size_t cur_degree ) const { const auto* batch_data = get_batch_data(data_id); @@ -378,7 +369,7 @@ void QuantizedGraph::scan_neighbors( template inline void QuantizedGraph::update_results( - buffer::SearchBuffer& result_pool, HashBasedBooleanSet& vis, const T* query + buffer::SearchBuffer& result_pool, VisitedSet& vis, const T* query ) { if (result_pool.is_full()) { return; @@ -433,7 +424,7 @@ inline void QuantizedGraph::find_candidates( PID cur_id, size_t search_ef, std::vector>& results, - HashBasedBooleanSet& vis, + VisitedSet& vis, const std::vector& degrees ) const { const T* query = get_vector(cur_id); diff --git a/include/rabitqlib/index/symqg/qg_builder.hpp b/include/rabitqlib/index/symqg/qg_builder.hpp index 9d90248..27ffef3 100644 --- a/include/rabitqlib/index/symqg/qg_builder.hpp +++ b/include/rabitqlib/index/symqg/qg_builder.hpp @@ -11,9 +11,9 @@ #include "rabitqlib/defines.hpp" #include "rabitqlib/index/symqg/qg.hpp" -#include "rabitqlib/utils/hashset.hpp" #include "rabitqlib/utils/space.hpp" #include "rabitqlib/utils/tools.hpp" +#include "rabitqlib/utils/visited_set.hpp" namespace rabitqlib::symqg { constexpr size_t kMaxBsIter = 5; // max iter for binary search of pruning bar @@ -37,9 +37,9 @@ class QGBuilder { static constexpr size_t kMaxPrunedSize = 300; // max number of recorded pruned candidates std::vector new_neighbors_; // new neighbors for current iteration - std::vector pruned_neighbors_; // recorded pruned neighbors - std::vector visited_list_; // list of visited hash set - std::vector degrees_; // record degree of qg + std::vector pruned_neighbors_; // recorded pruned neighbors + std::vector visited_list_; // per-thread visited sets + std::vector degrees_; // record degree of qg void random_init(); void search_new_neighbors(bool refine); void heuristic_prune(PID, CandidateList&, CandidateList&, bool); @@ -67,7 +67,7 @@ class QGBuilder { , pruned_neighbors_(qg_.num_vertices()) , visited_list_( num_threads_, - HashBasedBooleanSet(std::min(ef_build_ * ef_build_, num_nodes_ / 10)) + VisitedSet(num_nodes_, std::min(ef_build_ * ef_build_, num_nodes_ / 10)) ) , degrees_(qg_.num_vertices(), degree_bound_) { omp_set_num_threads(static_cast(num_threads_)); @@ -236,7 +236,7 @@ inline void QGBuilder::search_new_neighbors(bool refine) { PID cur_id = i; auto tid = omp_get_thread_num(); CandidateList candidates; - HashBasedBooleanSet& vis = visited_list_[tid]; + VisitedSet& vis = visited_list_[tid]; candidates.reserve(2 * kMaxCandidatePoolSize); vis.clear(); qg_.find_candidates(cur_id, ef_build_, candidates, vis, degrees_); diff --git a/include/rabitqlib/utils/hashset.hpp b/include/rabitqlib/utils/hashset.hpp index b363789..2633417 100644 --- a/include/rabitqlib/utils/hashset.hpp +++ b/include/rabitqlib/utils/hashset.hpp @@ -1,101 +1,8 @@ -// This code is modified based on NGT from Yahoo Japan -// https://github.com/yahoojapan/NGT -// -// Copyright (C) 2015 Yahoo Japan Corporation -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. -// - +// This compatibility header preserves the original visited-set API. #pragma once -#include -#include -#include -#include - -#include "rabitqlib/defines.hpp" -#include "rabitqlib/utils/memory.hpp" +#include "rabitqlib/utils/visited_set_hash.hpp" namespace rabitqlib { -/** - * @brief hash set to record visited vertices - * - */ -class HashBasedBooleanSet { - private: - size_t table_size_ = 0; - PID mask_ = 0; - std::vector> table_; - std::unordered_set stl_hash_; - - [[nodiscard]] auto hash1(const PID value) const { return value & mask_; } - - public: - HashBasedBooleanSet() = default; - ~HashBasedBooleanSet() = default; - - HashBasedBooleanSet(const HashBasedBooleanSet&) = default; - HashBasedBooleanSet(HashBasedBooleanSet&&) noexcept = default; - HashBasedBooleanSet& operator=(HashBasedBooleanSet&&) noexcept = default; - - explicit HashBasedBooleanSet(size_t size) { - size_t bit_size = 0; - size_t bit = size; - while (bit != 0) { - bit_size++; - bit >>= 1; - } - size_t bucket_size = 0x1 << ((bit_size + 4) / 2 + 3); - initialize(bucket_size); - } - - void initialize(const size_t table_size) { - table_size_ = table_size; - mask_ = static_cast(table_size_ - 1); - const PID check_val = hash1(static_cast(table_size)); - if (check_val != 0) { - std::cerr << "[WARN] table size is not 2^N : " << table_size << '\n'; - } - - table_ = std::vector>(table_size); - std::fill(table_.begin(), table_.end(), kPidMax); - stl_hash_.clear(); - } - - void clear() { - std::fill(table_.begin(), table_.end(), kPidMax); - stl_hash_.clear(); - } - - // get if data_id is in the hashset - [[nodiscard]] bool get(PID data_id) const { - PID val = this->table_[hash1(data_id)]; - if (val == data_id) { - return true; - } - return (val != kPidMax && stl_hash_.find(data_id) != stl_hash_.end()); - } - - void set(PID data_id) { - PID& val = table_[hash1(data_id)]; - if (val == data_id) { - return; - } - if (val == kPidMax) { - val = data_id; - } else { - stl_hash_.emplace(data_id); - } - } -}; -} // namespace rabitqlib \ No newline at end of file +using HashBasedBooleanSet = HashBasedVisitedSet; +} // namespace rabitqlib diff --git a/include/rabitqlib/utils/visited_pool.hpp b/include/rabitqlib/utils/visited_pool.hpp index ac9d27d..2926d3d 100644 --- a/include/rabitqlib/utils/visited_pool.hpp +++ b/include/rabitqlib/utils/visited_pool.hpp @@ -2,48 +2,48 @@ #include #include -#include "rabitqlib/utils/hashset.hpp" +#include "rabitqlib/utils/visited_set.hpp" namespace rabitqlib { class VisitedListPool { - std::deque pool_; + std::deque pool_; std::mutex poolguard_; size_t numelements_; public: VisitedListPool(size_t initpoolsize, size_t max_elements) { - numelements_ = max_elements / 10; + numelements_ = max_elements; for (size_t i = 0; i < initpoolsize; i++) { - pool_.push_front(new HashBasedBooleanSet(numelements_)); + pool_.push_front(new VisitedSet(numelements_, numelements_ / 10)); } } - HashBasedBooleanSet* get_free_vislist() { - HashBasedBooleanSet* rez; + VisitedSet* get_free_vislist() { + VisitedSet* rez; { std::unique_lock lock(poolguard_); if (pool_.size() > 0) { rez = pool_.front(); pool_.pop_front(); } else { - rez = new HashBasedBooleanSet(numelements_); + rez = new VisitedSet(numelements_, numelements_ / 10); } } rez->clear(); return rez; } - void release_vis_list(HashBasedBooleanSet* vl) { + void release_vis_list(VisitedSet* vl) { std::unique_lock lock(poolguard_); pool_.push_front(vl); } ~VisitedListPool() { while (pool_.size() > 0) { - HashBasedBooleanSet* rez = pool_.front(); + VisitedSet* rez = pool_.front(); pool_.pop_front(); ::delete rez; } } }; -} // namespace rabitqlib \ No newline at end of file +} // namespace rabitqlib diff --git a/include/rabitqlib/utils/visited_set.hpp b/include/rabitqlib/utils/visited_set.hpp new file mode 100644 index 0000000..6ca6c14 --- /dev/null +++ b/include/rabitqlib/utils/visited_set.hpp @@ -0,0 +1,46 @@ +#pragma once + +#include +#include + +#include "rabitqlib/defines.hpp" +#include "rabitqlib/utils/visited_set_epoch.hpp" +#include "rabitqlib/utils/visited_set_hash.hpp" + +namespace rabitqlib { +namespace detail { +template +struct CheckVisitedSet { + static_assert( + std::is_constructible::value, + "visited set must be constructible from a size_t id space" + ); + static_assert( + std::is_constructible::value, + "visited set must be constructible from (num_elements, size_hint)" + ); + static_assert( + std::is_same::value, + "visited set must declare: void initialize(size_t num_elements, size_t size_hint)" + ); + static_assert( + std::is_same::value, + "visited set must declare: void clear()" + ); + static_assert( + std::is_same::value, + "visited set must declare: bool get(PID) const" + ); + static_assert( + std::is_same::value, + "visited set must declare: void set(PID)" + ); + static constexpr bool kValue = true; +}; +} // namespace detail + +static_assert(detail::CheckVisitedSet::kValue, ""); +static_assert(detail::CheckVisitedSet::kValue, ""); + +using VisitedSet = HashBasedVisitedSet; +} // namespace rabitqlib diff --git a/include/rabitqlib/utils/visited_set_epoch.hpp b/include/rabitqlib/utils/visited_set_epoch.hpp new file mode 100644 index 0000000..04152a9 --- /dev/null +++ b/include/rabitqlib/utils/visited_set_epoch.hpp @@ -0,0 +1,46 @@ +#pragma once + +#include +#include +#include + +#include "rabitqlib/defines.hpp" +#include "rabitqlib/utils/memory.hpp" + +namespace rabitqlib { +class EpochBasedVisitedSet { + private: + std::vector> stamp_; + uint16_t cur_ = 0; + + public: + EpochBasedVisitedSet() = default; + ~EpochBasedVisitedSet() = default; + + EpochBasedVisitedSet(const EpochBasedVisitedSet&) = default; + EpochBasedVisitedSet& operator=(const EpochBasedVisitedSet&) = default; + EpochBasedVisitedSet(EpochBasedVisitedSet&&) noexcept = default; + EpochBasedVisitedSet& operator=(EpochBasedVisitedSet&&) noexcept = default; + + explicit EpochBasedVisitedSet(size_t num_elements, size_t = 0) { + initialize(num_elements); + } + + void initialize(size_t num_elements, size_t = 0) { + stamp_.assign(num_elements, 0); + cur_ = 0; + clear(); + } + + void clear() { + if (++cur_ == 0) { + std::fill(stamp_.begin(), stamp_.end(), 0); + cur_ = 1; + } + } + + [[nodiscard]] bool get(PID data_id) const { return stamp_[data_id] == cur_; } + + void set(PID data_id) { stamp_[data_id] = cur_; } +}; +} // namespace rabitqlib diff --git a/include/rabitqlib/utils/visited_set_hash.hpp b/include/rabitqlib/utils/visited_set_hash.hpp new file mode 100644 index 0000000..f9f3722 --- /dev/null +++ b/include/rabitqlib/utils/visited_set_hash.hpp @@ -0,0 +1,108 @@ +// This code is modified based on NGT from Yahoo Japan +// https://github.com/yahoojapan/NGT +// +// Copyright (C) 2015 Yahoo Japan Corporation +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/defines.hpp" +#include "rabitqlib/utils/memory.hpp" + +namespace rabitqlib { +class HashBasedVisitedSet { + public: + static constexpr size_t kNoSizeHint = std::numeric_limits::max(); + + private: + size_t table_size_ = 0; + PID mask_ = 0; + std::vector> table_; + std::unordered_set stl_hash_; + + [[nodiscard]] auto hash1(const PID value) const { return value & mask_; } + + void initialize_table(const size_t table_size) { + table_size_ = table_size; + mask_ = static_cast(table_size_ - 1); + const PID check_val = hash1(static_cast(table_size)); + if (check_val != 0) { + std::cerr << "[WARN] table size is not 2^N : " << table_size << '\n'; + } + + table_ = std::vector>(table_size); + std::fill(table_.begin(), table_.end(), kPidMax); + stl_hash_.clear(); + } + + public: + HashBasedVisitedSet() = default; + ~HashBasedVisitedSet() = default; + + HashBasedVisitedSet(const HashBasedVisitedSet&) = default; + HashBasedVisitedSet& operator=(const HashBasedVisitedSet&) = default; + HashBasedVisitedSet(HashBasedVisitedSet&&) noexcept = default; + HashBasedVisitedSet& operator=(HashBasedVisitedSet&&) noexcept = default; + + explicit HashBasedVisitedSet(size_t num_elements, size_t size_hint = kNoSizeHint) { + initialize(num_elements, size_hint); + } + + void initialize(size_t num_elements, size_t size_hint = kNoSizeHint) { + size_t size = std::min(num_elements, size_hint); + size_t bit_size = 0; + size_t bit = size; + while (bit != 0) { + bit_size++; + bit >>= 1; + } + size_t bucket_size = static_cast(0x1) << ((bit_size + 4) / 2 + 3); + initialize_table(bucket_size); + } + + void clear() { + std::fill(table_.begin(), table_.end(), kPidMax); + stl_hash_.clear(); + } + + [[nodiscard]] bool get(PID data_id) const { + PID val = this->table_[hash1(data_id)]; + if (val == data_id) { + return true; + } + return (val != kPidMax && stl_hash_.find(data_id) != stl_hash_.end()); + } + + void set(PID data_id) { + PID& val = table_[hash1(data_id)]; + if (val == data_id) { + return; + } + if (val == kPidMax) { + val = data_id; + } else { + stl_hash_.emplace(data_id); + } + } +}; +} // namespace rabitqlib diff --git a/tests/unit/rabitqlib/utils/visited_set_test.cpp b/tests/unit/rabitqlib/utils/visited_set_test.cpp new file mode 100644 index 0000000..61573b0 --- /dev/null +++ b/tests/unit/rabitqlib/utils/visited_set_test.cpp @@ -0,0 +1,138 @@ +#include "rabitqlib/utils/visited_set.hpp" + +#include + +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/utils/hashset.hpp" + +namespace { + +using rabitqlib::EpochBasedVisitedSet; +using rabitqlib::HashBasedBooleanSet; +using rabitqlib::HashBasedVisitedSet; +using rabitqlib::PID; +using rabitqlib::VisitedSet; + +static_assert( + std::is_same_v, + "HashBasedVisitedSet must remain the default visited-set implementation" +); +static_assert( + std::is_same_v, + "HashBasedBooleanSet must remain available as a compatibility alias" +); + +template +void expect_basic_visited_set_behavior() { + constexpr size_t kNumElements = 1024; + Set visited(kNumElements, 8); + + EXPECT_FALSE(visited.get(0)); + EXPECT_FALSE(visited.get(kNumElements - 1)); + + visited.set(0); + visited.set(17); + visited.set(kNumElements - 1); + + EXPECT_TRUE(visited.get(0)); + EXPECT_TRUE(visited.get(17)); + EXPECT_TRUE(visited.get(kNumElements - 1)); + + visited.set(17); + EXPECT_TRUE(visited.get(17)); + + visited.clear(); + EXPECT_FALSE(visited.get(0)); + EXPECT_FALSE(visited.get(17)); + EXPECT_FALSE(visited.get(kNumElements - 1)); + + visited.set(23); + EXPECT_TRUE(visited.get(23)); +} + +TEST(VisitedSet, HashBackendSatisfiesContract) { + expect_basic_visited_set_behavior(); +} + +TEST(VisitedSet, EpochBackendSatisfiesContract) { + expect_basic_visited_set_behavior(); +} + +TEST(VisitedSet, BackendsRemainBehaviorallyEquivalent) { + constexpr size_t kNumElements = 4096; + constexpr size_t kNumOperations = 20000; + + HashBasedVisitedSet hash_visited(kNumElements, 16); + EpochBasedVisitedSet epoch_visited(kNumElements); + std::vector expected(kNumElements, false); + std::mt19937 generator(42); + std::uniform_int_distribution id_distribution(0, kNumElements - 1); + + for (size_t operation = 0; operation < kNumOperations; ++operation) { + if (operation % 257 == 0) { + hash_visited.clear(); + epoch_visited.clear(); + std::fill(expected.begin(), expected.end(), false); + } else { + const PID id = id_distribution(generator); + if (operation % 3 == 0) { + EXPECT_EQ(hash_visited.get(id), expected[id]); + EXPECT_EQ(epoch_visited.get(id), expected[id]); + } else { + hash_visited.set(id); + epoch_visited.set(id); + expected[id] = true; + } + } + } + + for (PID id = 0; id < kNumElements; ++id) { + EXPECT_EQ(hash_visited.get(id), expected[id]) << "id=" << id; + EXPECT_EQ(epoch_visited.get(id), expected[id]) << "id=" << id; + } +} + +TEST(HashBasedVisitedSet, HandlesCollidingIds) { + HashBasedVisitedSet visited(1024, 1); + const std::vector colliding_ids{0, 32, 64, 96, 128, 1023}; + + for (PID id : colliding_ids) { + visited.set(id); + } + for (PID id : colliding_ids) { + EXPECT_TRUE(visited.get(id)) << "id=" << id; + } + + visited.clear(); + for (PID id : colliding_ids) { + EXPECT_FALSE(visited.get(id)) << "id=" << id; + } +} + +TEST(EpochBasedVisitedSet, ClearsStaleValuesWhenEpochWraps) { + EpochBasedVisitedSet visited(4); + visited.set(0); + + // Construction starts at epoch 1. Advance to the last uint16_t epoch. + for (size_t epoch = 0; epoch < UINT16_MAX - 1; ++epoch) { + visited.clear(); + } + visited.set(1); + EXPECT_TRUE(visited.get(1)); + + // The next clear wraps the counter and must reset all stored stamps. + visited.clear(); + EXPECT_FALSE(visited.get(0)); + EXPECT_FALSE(visited.get(1)); + + visited.set(3); + EXPECT_TRUE(visited.get(3)); +} + +} // namespace