From 1e0bfac1edb1e0f5832514b2a70eab8a684320eb Mon Sep 17 00:00:00 2001
From: GabTux <hevrgabr@fit.cvut.cz>
Date: Sat, 23 Mar 2024 23:25:00 +0100
Subject: [PATCH] Thread API

---
 include/ppqsort.h                              | 10 ++++++++++
 include/ppqsort/parallel/cpp/mainloop_par.h    | 17 +++++++++++++----
 include/ppqsort/parallel/openmp/mainloop_par.h | 12 ++++++++++--
 3 files changed, 33 insertions(+), 6 deletions(-)

diff --git a/include/ppqsort.h b/include/ppqsort.h
index 070a60d..c802c9a 100644
--- a/include/ppqsort.h
+++ b/include/ppqsort.h
@@ -64,6 +64,11 @@ namespace ppqsort {
        impl::call_sort(std::forward<ExecutionPolicy>(policy), begin, end);
    }
 
+   template <typename ExecutionPolicy, typename RandomIt>
+   void sort(ExecutionPolicy&& policy, RandomIt begin, RandomIt end, const int threads) {
+       impl::call_sort(std::forward<ExecutionPolicy>(policy), begin, end, threads);
+   }
+
    template <typename RandomIt, typename Compare>
    void sort(RandomIt begin, RandomIt end, Compare comp) {
        impl::seq_ppqsort(begin, end, comp);
@@ -73,4 +78,9 @@ namespace ppqsort {
    void sort(ExecutionPolicy&& policy, RandomIt begin, RandomIt end, Compare comp) {
        impl::call_sort(std::forward<ExecutionPolicy>(policy), begin, end, comp);
    }
+
+   template <typename ExecutionPolicy, typename RandomIt, typename Compare>
+   void sort(ExecutionPolicy&& policy, RandomIt begin, RandomIt end, Compare comp, const int threads) {
+       impl::call_sort(std::forward<ExecutionPolicy>(policy), begin, end, comp, threads);
+   }
 }
\ No newline at end of file
diff --git a/include/ppqsort/parallel/cpp/mainloop_par.h b/include/ppqsort/parallel/cpp/mainloop_par.h
index 1476a11..50b40c8 100644
--- a/include/ppqsort/parallel/cpp/mainloop_par.h
+++ b/include/ppqsort/parallel/cpp/mainloop_par.h
@@ -18,6 +18,9 @@ namespace ppqsort::impl {
     namespace cpp {
 
     struct ThreadPools {
+        ThreadPools() = default;
+        explicit ThreadPools(const int threads) : partition(threads), tasks(threads) {}
+
         ThreadPool<> partition;
         ThreadPool<> tasks;
     };
@@ -127,16 +130,15 @@ namespace ppqsort::impl {
         }
     }
 
-
     template <bool Force_branchless = false,
              typename RandomIt,
              typename Compare = std::less<typename std::iterator_traits<RandomIt>::value_type>,
              bool Branchless = use_branchless<typename std::iterator_traits<RandomIt>::value_type, Compare>::value>
-    void par_ppqsort(RandomIt begin, RandomIt end, Compare comp = Compare()) {
+    void par_ppqsort(RandomIt begin, RandomIt end, Compare comp = Compare(),
+                     int threads = static_cast<int>(std::jthread::hardware_concurrency())) {
         if (begin == end)
             return;
         constexpr bool branchless = Force_branchless || Branchless;
-        int threads = static_cast<int>(std::jthread::hardware_concurrency());
         auto size = end - begin;
         if ((threads < 2) || (size < parameters::seq_threshold))
             return seq_loop<RandomIt, Compare, branchless>(begin, end, comp, log2(size));
@@ -144,7 +146,7 @@ namespace ppqsort::impl {
         int seq_thr = (end - begin + 1) / threads / parameters::par_thr_div;
         seq_thr = std::max(seq_thr, branchless ? parameters::insertion_threshold_primitive
                                                    : parameters::insertion_threshold);
-        cpp::ThreadPools threadpools;
+        cpp::ThreadPools threadpools(threads);
         threadpools.tasks.push_task([begin, end, comp, seq_thr, threads, &threadpools] {
             cpp::par_loop<RandomIt, Compare, branchless>(begin, end, comp,
                       log2(end - begin),
@@ -154,4 +156,11 @@ namespace ppqsort::impl {
         threadpools.tasks.wait_and_stop();
         threadpools.partition.wait_and_stop();
     }
+
+    template <bool Force_branchless = false,
+         typename RandomIt,
+         typename Compare = std::less<typename std::iterator_traits<RandomIt>::value_type>>
+    void par_ppqsort(RandomIt begin, RandomIt end, int threads) {
+        return par_ppqsort<Force_branchless>(begin, end, Compare(), threads);
+    }
 }
\ No newline at end of file
diff --git a/include/ppqsort/parallel/openmp/mainloop_par.h b/include/ppqsort/parallel/openmp/mainloop_par.h
index 38095e5..afb7454 100644
--- a/include/ppqsort/parallel/openmp/mainloop_par.h
+++ b/include/ppqsort/parallel/openmp/mainloop_par.h
@@ -114,11 +114,11 @@ namespace ppqsort::impl {
          typename RandomIt,
          typename Compare = std::less<typename std::iterator_traits<RandomIt>::value_type>,
          bool Branchless = use_branchless<typename std::iterator_traits<RandomIt>::value_type, Compare>::value>
-    void par_ppqsort(RandomIt begin, RandomIt end, Compare comp = Compare()) {
+    void par_ppqsort(RandomIt begin, RandomIt end,
+                     Compare comp = Compare(), const int threads = omp_get_max_threads()) {
         if (begin == end)
             return;
         constexpr bool branchless = Force_branchless || Branchless;
-        const int threads = omp_get_max_threads();
         auto size = end - begin;
         if ((threads < 2) || (size < parameters::seq_threshold))
             return seq_loop<RandomIt, Compare, branchless>(begin, end, comp, log2(size));
@@ -138,4 +138,12 @@ namespace ppqsort::impl {
             }
         }
     }
+
+    template <bool Force_branchless = false,
+             typename RandomIt,
+             typename Compare = std::less<typename std::iterator_traits<RandomIt>::value_type>>
+        void par_ppqsort(RandomIt begin, RandomIt end, const int threads) {
+            Compare comp = Compare();
+            par_ppqsort<Force_branchless>(begin, end, comp, threads);
+        }
 }
\ No newline at end of file