fixed parallel_for for generic step logic, increased worker chase-lev queue size, minor formatting changes

Author: AlePalu

fdaPDE/src/execution/concurrency.h

@@ -51,8 +51,8 @@ inline std::size_t available_concurrency() noexcept {
     // SLURM
     if (auto SLURM_NCPUS = internals::get_env_concurrency_count("SLURM_CPUS_ON_NODE")) { return *SLURM_NCPUS; }
     // check if OpenMP has been configured
-    if (auto OMP_NUM_THREADS = internals::get_env_concurrency_count("OMP_NUM_THREADS")) {
-        if (*OMP_NUM_THREADS > 1) return *OMP_NUM_THREADS;
+    if (auto OMP_NCPUS   = internals::get_env_concurrency_count("OMP_NUM_THREADS")) {
+        if (*OMP_NCPUS > 1) return *OMP_NCPUS;
     }
 
     // second, check OS-specific settings

fdaPDE/src/execution/parallel_for.h

@@ -26,44 +26,66 @@ namespace internals {
 struct task_parallel_for {
     task_parallel_for() = default;
 
-    template <typename LoopBody, typename NextFunctor>
-        requires(std::is_invocable_v<LoopBody, int> && std::is_invocable_r_v<int, NextFunctor, int>)
-    void run(threaded_executor_impl* executor, int begin, int end, int grain_size, LoopBody&& f, NextFunctor&& next) {
-        int n = 0;
-        for (int i = begin; i < end; i = next(i), n++);
-        if (n <= 0) return;   // nothing to loop on
-
-        grain_size = std::max(1, std::min(grain_size, n));
-        dispatch_(executor, begin, end, grain_size, std::forward<LoopBody>(f), std::forward<NextFunctor>(next));
-        return;
-    }
+    // optimized for loop with standard ++i increment
     template <typename LoopBody>
         requires(std::is_invocable_v<LoopBody, int>)
     void run(threaded_executor_impl* executor, int begin, int end, int grain_size, LoopBody&& f) {
-        const int n = end - begin;
-        if (n <= 0) return;   // nothing to loop on
+        const int size = end - begin;
+        if (size <= 0) return;   // nothing to loop on
+
+        grain_size = std::max(1, std::min(grain_size, size));
+        std::atomic<int> local_task_count {1};
+
+        for (int local_begin = begin; local_begin < end; local_begin += grain_size) {
+            // register task in the group
+            local_task_count.fetch_add(1, std::memory_order_release);
 
-        grain_size = std::max(1, std::min(grain_size, n));
-        dispatch_(executor, begin, end, grain_size, std::forward<LoopBody>(f), [](int i) { return ++i; });
+            int local_end = ((end - local_begin) < grain_size) ? end : (local_begin + grain_size);
+            auto loop_body = [local_begin, local_end, &f, &local_task_count]() {
+                for (int it = local_begin; it < local_end; ++it) { f(it); }
+                // signal task completion
+                local_task_count.fetch_sub(1, std::memory_order_release);
+            };
+            executor->execute(std::move(loop_body));
+        }
+        // task_group collaborative wait
+        local_task_count.fetch_sub(1, std::memory_order_release);
+        executor->active_join(this_thread_id(), [&] {
+            // help the pool while the task group is not fully consumed
+            return local_task_count.load(std::memory_order_acquire) > 0;
+        });
         return;
     }
-   private:
+    // for loop with custom step logic
     template <typename LoopBody, typename NextFunctor>
-    void
-    dispatch_(threaded_executor_impl* executor, int begin, int end, int grain_size, LoopBody&& f, NextFunctor&& next) {
+        requires(std::is_invocable_v<LoopBody, int> && std::is_invocable_r_v<int, NextFunctor, int>)
+    void run(threaded_executor_impl* executor, int begin, int end, int grain_size, LoopBody&& f, NextFunctor&& next) {
+        int size = 0;
+        for (int i = begin; i < end; i = next(i), size++);
+        if (size <= 0) return;   // nothing to loop on
+
+        grain_size = std::max(1, std::min(grain_size, size));
         std::atomic<int> local_task_count {1};
+        int n_batches = std::ceil(double(size) / grain_size);
 
-        for (int j = begin; j < end; j += grain_size) {
+        int local_begin = begin;
+        int local_end = local_begin;
+        for (int i = 0; i < grain_size && local_end < end; ++i) { local_end = next(local_end); }
+
+        for (int j = 0; j < n_batches; j++) {
             // register task in the group
             local_task_count.fetch_add(1, std::memory_order_release);
-
-            int k = ((end - j) < grain_size) ? end : (j + grain_size);
-            auto loop_body = [j, k, next, &f, &local_task_count]() {
-                for (int it = j; it < k; it = next(it)) { f(it); }
+            auto loop_body = [local_begin, local_end, next, &f, &local_task_count]() {
+                for (int it = local_begin; it < local_end; it = next(it)) { f(it); }
                 // signal task completion
                 local_task_count.fetch_sub(1, std::memory_order_release);
             };
             executor->execute(std::move(loop_body));
+
+            // update next task range
+            local_begin = local_end;
+            local_end = local_begin;
+            for (int i = 0; i < grain_size && local_end < end; ++i) { local_end = next(local_end); }
         }
         // task_group collaborative wait
         local_task_count.fetch_sub(1, std::memory_order_release);

fdaPDE/src/execution/worker.h

@@ -314,7 +314,7 @@ struct worker {
     using task_const_pointer = typename allocator_type::const_pointer;
     using task_queue_type = chase_lev_queue<task_pointer>;
     using task_buffer_type = mpsc_queue<task_pointer>;
-    static constexpr int task_queue_size = 4096;
+    static constexpr int task_queue_size = 8192;   // 2^13
 
     // constructor
     template <typename Executor>