Add a pushLoop method to ThreadPool

include/slang/util/ThreadPool.h

@@ -60,6 +60,7 @@ class ThreadPool {
             thread.join();
     }
 
+    /// Gets the number of threads in the thread pool.
     size_t getThreadCount() const { return threads.size(); }
 
     /// @brief Pushes a new task into the pool for execution.
@@ -107,6 +108,35 @@ class ThreadPool {
         return taskPromise->get_future();
     }
 
+    /// @brief Pushes several tasks into the pool in order to parallelize
+    /// the loop given by [from, to).
+    ///
+    /// The loop will be broken into a number of blocks as specified by
+    /// @a numBlocks -- or if zero, the number of blocks will be set to
+    /// the number of threads in the pool.
+    template<typename TIndex, typename TFunc>
+    void pushLoop(TIndex from, TIndex to, TFunc&& body, size_t numBlocks = 0) {
+        SLANG_ASSERT(to >= from);
+        if (!numBlocks)
+            numBlocks = getThreadCount();
+
+        const size_t totalSize = size_t(to - from);
+        if (!totalSize)
+            return;
+
+        size_t blockSize = totalSize / numBlocks;
+        if (blockSize == 0) {
+            blockSize = 1;
+            numBlocks = totalSize;
+        }
+
+        for (size_t i = 0; i < numBlocks; i++) {
+            const TIndex start = TIndex(i * blockSize) + from;
+            const TIndex end = i == numBlocks - 1 ? to : TIndex(start + blockSize);
+            pushTask(std::forward<TFunc>(body), start, end);
+        }
+    }
+
     /// Blocks the calling thread until all running tasks are complete.
     void waitForAll() {
         std::unique_lock lock(mutex);

tests/unittests/util/ThreadPoolTests.cpp

@@ -235,6 +235,34 @@ TEST_CASE("ThreadPool -- exception handling") {
 }
 #endif
 
+TEST_CASE("ThreadPool -- pushLoop") {
+    ThreadPool pool(3);
+
+    // Pushing an empty loop does nothing.
+    bool flag = false;
+    pool.pushLoop(3, 3, [&](int start, int end) { flag = true; });
+    pool.waitForAll();
+    CHECK(!flag);
+
+    std::array<std::atomic<bool>, 2> flags2{false, false};
+    pool.pushLoop(3, 5, [&](int start, int end) {
+        for (int i = start; i < end; i++)
+            flags2[i - 3] = true;
+    });
+    pool.waitForAll();
+    CHECK(std::ranges::all_of(flags2, [](auto&& f) -> bool { return f; }));
+
+    std::array<std::atomic<bool>, 10> flags10;
+    std::ranges::fill(flags10, false);
+
+    pool.pushLoop(3, 13, [&](int start, int end) {
+        for (int i = start; i < end; i++)
+            flags10[i - 3] = true;
+    });
+    pool.waitForAll();
+    CHECK(std::ranges::all_of(flags10, [](auto&& f) -> bool { return f; }));
+}
+
 #ifdef CI_BUILD
 
 TEST_CASE("ThreadPool -- no destruction deadlocks") {