Add a ForkJoin-based ParallelFor algorithm to the ThreadPool module

test/threads_fork_join.cpp

@@ -0,0 +1,79 @@
+// This file is part of Eigen, a lightweight C++ template library
+// for linear algebra.
+//
+// Copyright (C) 2025 Weiwei Kong <weiweikong@google.com>
+//
+// This Source Code Form is subject to the terms of the Mozilla
+// Public License v. 2.0. If a copy of the MPL was not distributed
+// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+#define EIGEN_USE_THREADS
+#include "main.h"
+#include "Eigen/ThreadPool"
+
+struct TestData {
+  ThreadPool tp;
+  std::vector<double> data;
+};
+
+TestData make_test_data(int num_threads, int num_shards) {
+  return {ThreadPool(num_threads), std::vector<double>(num_shards, 1.0)};
+}
+
+static void test_unary_parallel_for(int granularity) {
+  // Test correctness.
+  const int kNumTasks = 100000;
+  TestData test_data = make_test_data(/*num_threads=*/4, kNumTasks);
+  std::atomic<double> sum = 0.0;
+  std::function<void(int)> unary_do_fn = [&](int i) {
+    for (double new_sum = sum; !sum.compare_exchange_weak(new_sum, new_sum + test_data.data[i]);) {
+    };
+  };
+  ForkJoinScheduler::ParallelFor(0, kNumTasks, granularity, std::move(unary_do_fn), &test_data.tp);
+  VERIFY_IS_EQUAL(sum, kNumTasks);
+}
+
+static void test_binary_parallel_for(int granularity) {
+  // Test correctness.
+  const int kNumTasks = 100000;
+  TestData test_data = make_test_data(/*num_threads=*/4, kNumTasks);
+  std::atomic<double> sum = 0.0;
+  std::function<void(int, int)> binary_do_fn = [&](int i, int j) {
+    for (int k = i; k < j; ++k)
+      for (double new_sum = sum; !sum.compare_exchange_weak(new_sum, new_sum + test_data.data[k]);) {
+      };
+  };
+  ForkJoinScheduler::ParallelFor(0, kNumTasks, granularity, std::move(binary_do_fn), &test_data.tp);
+  VERIFY_IS_EQUAL(sum, kNumTasks);
+}
+
+static void test_async_parallel_for() {
+  // Test correctness.
+  // NOTE: Granularity and type of `do_func` are checked in the synchronous tests.
+  const int kNumThreads = 4;
+  const int kNumTasks = 100;
+  const int kNumAsyncCalls = kNumThreads * 4;
+  TestData test_data = make_test_data(kNumThreads, kNumTasks);
+  std::atomic<double> sum = 0.0;
+  std::function<void(int)> unary_do_fn = [&](int i) {
+    for (double new_sum = sum; !sum.compare_exchange_weak(new_sum, new_sum + test_data.data[i]);) {
+    };
+  };
+  Barrier barrier(kNumTasks * kNumAsyncCalls);
+  std::function<void()> done = [&]() { barrier.Notify(); };
+  for (int k = 0; k < kNumAsyncCalls; ++k) {
+    test_data.tp.Schedule([&]() {
+      ForkJoinScheduler::ParallelForAsync(0, kNumTasks, /*granularity=*/1, unary_do_fn, done, &test_data.tp);
+    });
+  }
+  barrier.Wait();
+  VERIFY_IS_EQUAL(sum, kNumTasks * kNumAsyncCalls);
+}
+
+EIGEN_DECLARE_TEST(fork_join) {
+  CALL_SUBTEST(test_unary_parallel_for(1));
+  CALL_SUBTEST(test_unary_parallel_for(2));
+  CALL_SUBTEST(test_binary_parallel_for(1));
+  CALL_SUBTEST(test_binary_parallel_for(2));
+  CALL_SUBTEST(test_async_parallel_for());
+}