Enable parallelization for insert+search in HNSW - MOD-4628 (#316)

alonre24 · web-flow · commit 32ebe982c854 · 2023-01-11T18:37:35.000+02:00
diff --git a/src/VecSim/algorithms/hnsw/hnsw.h b/src/VecSim/algorithms/hnsw/hnsw.h
diff --git a/src/VecSim/algorithms/hnsw/hnsw_batch_iterator.h b/src/VecSim/algorithms/hnsw/hnsw_batch_iterator.h
@@ -78,7 +78,7 @@ HNSW_BatchIterator<DataType, DistType>::HNSW_BatchIterator(
 
     this->dist_func = index->getDistFunc();
     this->dim = index->getDim();
-    this->entry_point = index->getEntryPointId();
+    this->entry_point = HNSW_INVALID_ID; // temporary until we store the entry point to level 0.
     // Use "fresh" tag to mark nodes that were visited along the search in some iteration.
     this->visited_list = index->getVisitedList();
     this->visited_tag = this->visited_list->getFreshTag();
@@ -114,6 +114,7 @@ VecSimQueryResult_Code HNSW_BatchIterator<DataType, DistType>::scanGraphInternal
 
         // Take the current node out of the candidates queue and go over his neighbours.
         candidates.pop();
+        this->index->lockNodeLinks(curr_node_id);
         idType *node_links = this->index->get_linklist_at_level(curr_node_id, 0);
         linkListSize links_num = this->index->getListCount(node_links);
 
@@ -137,6 +138,7 @@ VecSimQueryResult_Code HNSW_BatchIterator<DataType, DistType>::scanGraphInternal
             candidates.emplace(candidate_dist, candidate_id);
             __builtin_prefetch(index->get_linklist_at_level(candidates.top().second, 0));
         }
+        this->index->unlockNodeLinks(curr_node_id);
     }
     return VecSim_QueryResult_OK;
 }
diff --git a/src/VecSim/algorithms/hnsw/hnsw_factory.cpp b/src/VecSim/algorithms/hnsw/hnsw_factory.cpp
@@ -63,6 +63,7 @@ size_t EstimateInitialSize(const HNSWParams *params) {
         est += sizeof(size_t) * params->initialCapacity + sizeof(size_t); // element level
         est += sizeof(size_t) * params->initialCapacity +
                sizeof(size_t); // Labels lookup hash table buckets.
+        est += sizeof(std::mutex) * params->initialCapacity + sizeof(size_t); // lock per vector
     }
 
     // Explicit allocation calls - always allocate a header.
@@ -115,6 +116,7 @@ size_t EstimateElementSize(const HNSWParams *params) {
     // lookup hash map.
     size_t size_meta_data =
         sizeof(tag_t) + sizeof(size_t) + sizeof(size_t) + size_label_lookup_node;
+    size_t size_lock = sizeof(std::mutex);
 
     /* Disclaimer: we are neglecting two additional factors that consume memory:
      * 1. The overall bucket size in labels_lookup hash table is usually higher than the number of
@@ -123,7 +125,7 @@ size_t EstimateElementSize(const HNSWParams *params) {
      * 2. The incoming edges that aren't bidirectional are stored in a dynamic array
      * (vecsim_stl::vector) Those edges' memory *is omitted completely* from this estimation.
      */
-    return size_meta_data + size_total_data_per_element;
+    return size_meta_data + size_total_data_per_element + size_lock;
 }
 
 VecSimIndex *NewTieredIndex(const TieredHNSWParams *params,
diff --git a/src/VecSim/algorithms/hnsw/hnsw_multi.h b/src/VecSim/algorithms/hnsw/hnsw_multi.h
@@ -73,6 +73,8 @@ class HNSWIndex_Multi : public HNSWIndex<DataType, DistType> {
     int addVector(const void *vector_data, labelType label, bool overwrite_allowed = true) override;
     double getDistanceFrom(labelType label, const void *vector_data) const override;
     inline std::vector<idType> markDelete(labelType label) override;
+    inline bool safeCheckIfLabelExistsInIndex(labelType label,
+                                              bool also_done_processing) const override;
 };
 
 /**
@@ -188,3 +190,23 @@ std::vector<idType> HNSWIndex_Multi<DataType, DistType>::markDelete(labelType la
     label_lookup_.erase(search);
     return idsToDelete;
 }
+
+template <typename DataType, typename DistType>
+inline bool HNSWIndex_Multi<DataType, DistType>::safeCheckIfLabelExistsInIndex(
+    labelType label, bool also_done_processing) const {
+    std::unique_lock<std::mutex> index_data_lock(this->index_data_guard_);
+    auto search_res = label_lookup_.find(label);
+    bool exists = search_res != label_lookup_.end();
+    // If we want to make sure that the vector(s) stored under the label were already indexed,
+    // we go on and check that every associated vector is no longer in process.
+    if (exists && also_done_processing) {
+        for (auto id : search_res->second) {
+            exists = !this->isInProcess(id);
+            // If we find at least one internal id that is still in process, consider it as not
+            // ready.
+            if (!exists)
+                return false;
+        }
+    }
+    return exists;
+}
diff --git a/src/VecSim/algorithms/hnsw/hnsw_serializer.h b/src/VecSim/algorithms/hnsw/hnsw_serializer.h
@@ -8,7 +8,8 @@ HNSWIndex<DataType, DistType>::HNSWIndex(std::ifstream &input, const HNSWParams
                                     params->blockSize, params->multi),
       Serializer(version), max_elements_(params->initialCapacity), epsilon_(params->epsilon),
       element_levels_(max_elements_, allocator),
-      visited_nodes_handler_pool(1, max_elements_, allocator) {
+      visited_nodes_handler_pool(1, max_elements_, allocator),
+      element_neighbors_locks_(max_elements_, allocator) {
 
     this->restoreIndexFields(input);
     this->fieldsValidation();
diff --git a/src/VecSim/algorithms/hnsw/hnsw_single.h b/src/VecSim/algorithms/hnsw/hnsw_single.h
@@ -64,6 +64,8 @@ class HNSWIndex_Single : public HNSWIndex<DataType, DistType> {
     int addVector(const void *vector_data, labelType label, bool overwrite_allowed = true) override;
     double getDistanceFrom(labelType label, const void *vector_data) const override;
     inline std::vector<idType> markDelete(labelType label) override;
+    inline bool safeCheckIfLabelExistsInIndex(labelType label,
+                                              bool also_done_processing = false) const override;
 };
 
 /**
@@ -121,6 +123,7 @@ int HNSWIndex_Single<DataType, DistType>::addVector(const void *vector_data, con
                                                     bool overwrite_allowed) {
 
     // Checking if an element with the given label already exists.
+    std::unique_lock<std::mutex> index_data_lock(this->index_data_guard_);
     bool label_exists = false;
     if (label_lookup_.find(label) != label_lookup_.end()) {
         label_exists = true;
@@ -132,7 +135,7 @@ int HNSWIndex_Single<DataType, DistType>::addVector(const void *vector_data, con
             return -1;
         }
     }
-
+    index_data_lock.unlock();
     this->appendVector(vector_data, label);
     // Return the delta in the index size due to the insertion.
     return label_exists ? 0 : 1;
@@ -159,6 +162,7 @@ HNSWIndex_Single<DataType, DistType>::newBatchIterator(const void *queryBlob,
 template <typename DataType, typename DistType>
 std::vector<idType> HNSWIndex_Single<DataType, DistType>::markDelete(labelType label) {
     std::vector<idType> idsToDelete;
+    std::unique_lock<std::mutex> index_data_lock(this->index_data_guard_);
     auto search = label_lookup_.find(label);
     if (search == label_lookup_.end()) {
         return idsToDelete;
@@ -168,3 +172,17 @@ std::vector<idType> HNSWIndex_Single<DataType, DistType>::markDelete(labelType l
     label_lookup_.erase(search);
     return idsToDelete;
 }
+
+template <typename DataType, typename DistType>
+inline bool HNSWIndex_Single<DataType, DistType>::safeCheckIfLabelExistsInIndex(
+    labelType label, bool also_done_processing) const {
+    std::unique_lock<std::mutex> index_data_lock(this->index_data_guard_);
+    auto it = label_lookup_.find(label);
+    bool exists = it != label_lookup_.end();
+    // If we want to make sure that the vector stored under the label was already indexed,
+    // we go on and check that its associated internal id is no longer in process.
+    if (exists && also_done_processing) {
+        return !this->isInProcess(it->second);
+    }
+    return exists;
+}
diff --git a/src/VecSim/algorithms/hnsw/hnsw_single_tests_friends.h b/src/VecSim/algorithms/hnsw/hnsw_single_tests_friends.h
@@ -10,3 +10,4 @@ INDEX_TEST_FRIEND_CLASS(HNSWTest_preferAdHocOptimization_Test)
 INDEX_TEST_FRIEND_CLASS(HNSWTest_testSizeEstimation_Test)
 INDEX_TEST_FRIEND_CLASS(IndexAllocatorTest_testIncomingEdgesSet_Test)
 INDEX_TEST_FRIEND_CLASS(IndexAllocatorTest_test_hnsw_reclaim_memory_Test)
+INDEX_TEST_FRIEND_CLASS(HNSWTestParallel_parallelInsertSearch_Test)
diff --git a/src/VecSim/algorithms/hnsw/visited_nodes_handler.cpp b/src/VecSim/algorithms/hnsw/visited_nodes_handler.cpp
@@ -41,11 +41,11 @@ VisitedNodesHandler::~VisitedNodesHandler() { allocator->free_allocation(element
 /**
  * VisitedNodesHandlerPool methods to enable parallel graph scans.
  */
-VisitedNodesHandlerPool::VisitedNodesHandlerPool(int initial_pool_size, int cap,
+VisitedNodesHandlerPool::VisitedNodesHandlerPool(size_t initial_pool_size, int cap,
                                                  const std::shared_ptr<VecSimAllocator> &allocator)
     : VecsimBaseObject(allocator), pool(initial_pool_size, allocator), num_elements(cap),
       total_handlers_in_use(1) {
-    for (int i = 0; i < initial_pool_size; i++)
+    for (size_t i = 0; i < initial_pool_size; i++)
         pool[i] = new (allocator) VisitedNodesHandler(cap, allocator);
 }
 
diff --git a/src/VecSim/algorithms/hnsw/visited_nodes_handler.h b/src/VecSim/algorithms/hnsw/visited_nodes_handler.h
@@ -58,7 +58,7 @@ class VisitedNodesHandlerPool : public VecsimBaseObject {
     unsigned short total_handlers_in_use;
 
 public:
-    VisitedNodesHandlerPool(int initial_pool_size, int cap,
+    VisitedNodesHandlerPool(size_t initial_pool_size, int cap,
                             const std::shared_ptr<VecSimAllocator> &allocator);
 
     VisitedNodesHandler *getAvailableVisitedNodesHandler();
diff --git a/src/VecSim/utils/vec_utils.h b/src/VecSim/utils/vec_utils.h
@@ -81,7 +81,7 @@ void normalizeVector(DataType *input_vector, size_t dim) {
     double sum = 0;
 
     for (size_t i = 0; i < dim; i++) {
-        sum += input_vector[i] * input_vector[i];
+        sum += (double)input_vector[i] * (double)input_vector[i];
     }
     DataType norm = sqrt(sum);
 
diff --git a/src/VecSim/vec_sim_common.h b/src/VecSim/vec_sim_common.h
@@ -33,9 +33,6 @@ typedef enum { VecSimAlgo_BF, VecSimAlgo_HNSWLIB } VecSimAlgo;
 // Distance metric
 typedef enum { VecSimMetric_L2, VecSimMetric_IP, VecSimMetric_Cosine } VecSimMetric;
 
-// Vectors flags (for marking a specific vector)
-typedef enum { DELETE_MARK = 0x01 } Flags;
-
 typedef size_t labelType;
 typedef unsigned int idType;
 
diff --git a/src/VecSim/vec_sim_interface.h b/src/VecSim/vec_sim_interface.h
@@ -43,7 +43,7 @@ struct VecSimIndexInterface : public VecsimBaseObject {
      * @return the number of new vectors inserted (1 for new insertion, 0 for override), or -1
      * in case that override is not allowed and label already exists.
      */
-    virtual int addVector(const void *blob, labelType label, bool overwriteAllowed = true) = 0;
+    virtual int addVector(const void *blob, labelType label, bool overwriteAllowed) = 0;
 
     /**
      * @brief Remove a vector from an index.
diff --git a/tests/unit/CMakeLists.txt b/tests/unit/CMakeLists.txt
@@ -22,12 +22,14 @@ set(CMAKE_EXE_LINKER_FLAGS "${CMAKE_LINKER_FLAGS} ${LLVM_LD_FLAGS}")
 enable_testing()
 
 add_executable(test_hnsw test_hnsw.cpp test_hnsw_multi.cpp test_hnsw_tiered.cpp test_utils.cpp)
+add_executable(test_hnsw_parallel test_hnsw_parallel.cpp test_utils.cpp)
 add_executable(test_bruteforce test_bruteforce.cpp test_bruteforce_multi.cpp test_utils.cpp)
 add_executable(test_allocator test_allocator.cpp test_utils.cpp)
 add_executable(test_spaces test_spaces.cpp)
 add_executable(test_common test_common.cpp test_utils.cpp)
 
 target_link_libraries(test_hnsw PUBLIC gtest_main VectorSimilarity)
+target_link_libraries(test_hnsw_parallel PUBLIC gtest_main VectorSimilarity)
 target_link_libraries(test_bruteforce PUBLIC gtest_main VectorSimilarity)
 target_link_libraries(test_allocator PUBLIC gtest_main VectorSimilarity)
 target_link_libraries(test_spaces PUBLIC gtest_main VectorSimilarity)
@@ -36,6 +38,7 @@ target_link_libraries(test_common PUBLIC gtest_main VectorSimilarity)
 include(GoogleTest)
 
 gtest_discover_tests(test_hnsw)
+gtest_discover_tests(test_hnsw_parallel)
 gtest_discover_tests(test_bruteforce)
 gtest_discover_tests(test_allocator)
 gtest_discover_tests(test_spaces)
diff --git a/tests/unit/test_allocator.cpp b/tests/unit/test_allocator.cpp
@@ -388,9 +388,9 @@ TYPED_TEST(IndexAllocatorTest, test_hnsw_reclaim_memory) {
     // Also account for all the memory allocation caused by the resizing that this vector triggered
     // except for the bucket count of the labels_lookup hash table that is calculated separately.
     size_t size_total_data_per_element = hnswIndex->size_data_per_element_;
-    expected_mem_delta +=
-        (sizeof(tag_t) + sizeof(void *) + sizeof(size_t) + size_total_data_per_element) *
-        block_size;
+    expected_mem_delta += (sizeof(tag_t) + sizeof(void *) + sizeof(size_t) +
+                           size_total_data_per_element + sizeof(std::mutex)) *
+                          block_size;
     expected_mem_delta +=
         (hnswIndex->label_lookup_.bucket_count() - prev_bucket_count) * sizeof(size_t);
 
diff --git a/tests/unit/test_hnsw.cpp b/tests/unit/test_hnsw.cpp
diff --git a/tests/unit/test_hnsw_multi.cpp b/tests/unit/test_hnsw_multi.cpp
diff --git a/tests/unit/test_hnsw_parallel.cpp b/tests/unit/test_hnsw_parallel.cpp

Original file line number	Diff line number	Diff line change
`@@ -81,7 +81,7 @@ void normalizeVector(DataType *input_vector, size_t dim) {`
`81`	`81`	`double sum = 0;`
`82`	`82`
`83`	`83`	`for (size_t i = 0; i < dim; i++) {`
`84`		`- sum += input_vector[i] * input_vector[i];`
	`84`	`+ sum += (double)input_vector[i] * (double)input_vector[i];`
`85`	`85`	`}`
`86`	`86`	`DataType norm = sqrt(sum);`
`87`	`87`