Parallel indexing

Author: alonre24

examples/searchKnnCloserFirst_test.cpp

@@ -9,12 +9,69 @@
 
 #include <vector>
 #include <iostream>
+#include <thread>
 
 namespace
 {
 
 using idx_t = hnswlib::labeltype;
 
+    template<class Function>
+    inline void ParallelFor(size_t start, size_t end, size_t numThreads, Function fn) {
+        if (numThreads <= 0) {
+            numThreads = std::thread::hardware_concurrency();
+        }
+
+        if (numThreads == 1) {
+            for (size_t id = start; id < end; id++) {
+                fn(id, 0);
+            }
+        } else {
+            std::vector<std::thread> threads;
+            std::atomic<size_t> current(start);
+
+            // keep track of exceptions in threads
+            // https://stackoverflow.com/a/32428427/1713196
+            std::exception_ptr lastException = nullptr;
+            std::mutex lastExceptMutex;
+
+            for (size_t threadId = 0; threadId < numThreads; ++threadId) {
+                threads.push_back(std::thread([&, threadId] {
+                    while (true) {
+                        size_t id = current.fetch_add(1);
+
+                        if ((id >= end)) {
+                            break;
+                        }
+
+                        try {
+                            fn(id, threadId);
+                        } catch (...) {
+                            std::unique_lock<std::mutex> lastExcepLock(lastExceptMutex);
+                            lastException = std::current_exception();
+                            /*
+                             * This will work even when current is the largest value that
+                             * size_t can fit, because fetch_add returns the previous value
+                             * before the increment (what will result in overflow
+                             * and produce 0 instead of current + 1).
+                             */
+                            current = end;
+                            break;
+                        }
+                    }
+                }));
+            }
+            for (auto &thread : threads) {
+                thread.join();
+            }
+            if (lastException) {
+                std::rethrow_exception(lastException);
+            }
+        }
+
+
+    }
+
 void test() {
     int d = 4;
     idx_t n = 100;
@@ -40,10 +97,18 @@ void test() {
     hnswlib::AlgorithmInterface<float>* alg_brute  = new hnswlib::BruteforceSearch<float>(&space, 2 * n);
     hnswlib::AlgorithmInterface<float>* alg_hnsw = new hnswlib::HierarchicalNSW<float>(&space, 2 * n);
 
-    for (size_t i = 0; i < n; ++i) {
-        alg_brute->addPoint(data.data() + d * i, i);
+//    for (size_t i = 0; i < n; ++i) {
+//        alg_brute->addPoint(data.data() + d * i, i);
+//        alg_hnsw->addPoint(data.data() + d * i, i);
+//    }
+
+    ParallelFor(0, n, 4, [&](size_t i, size_t threadId) {
         alg_hnsw->addPoint(data.data() + d * i, i);
-    }
+    });
+
+    ParallelFor(0, n, 4, [&](size_t i, size_t threadId) {
+        alg_brute->addPoint(data.data() + d * i, i);
+    });
 
     // test searchKnnCloserFirst of BruteforceSearch
     for (size_t j = 0; j < nq; ++j) {