Add parallel std::partition and use it to improve sort (#31)

This version of partition only uses two threads (spawns and waits for one additional task).

Author: alugowski

README.md

@@ -47,6 +47,7 @@ Algorithms are added on an as-needed basis. If you need one [open an issue](http
 * [`fill`](https://en.cppreference.com/w/cpp/algorithm/fill), [`fill_n`](https://en.cppreference.com/w/cpp/algorithm/fill_n)
 * [`find`](https://en.cppreference.com/w/cpp/algorithm/find), [`find_if`](https://en.cppreference.com/w/cpp/algorithm/find_if), [`find_if_not`](https://en.cppreference.com/w/cpp/algorithm/find_if_not)
 * [`for_each`](https://en.cppreference.com/w/cpp/algorithm/for_each), [`for_each_n`](https://en.cppreference.com/w/cpp/algorithm/for_each_n)
+* [`partition`](https://en.cppreference.com/w/cpp/algorithm/partition)
 * [`sort`](https://en.cppreference.com/w/cpp/algorithm/sort), [`stable_sort`](https://en.cppreference.com/w/cpp/algorithm/stable_sort)
 * [`transform`](https://en.cppreference.com/w/cpp/algorithm/transform)
 
@@ -197,9 +198,9 @@ for_each()/real_time                                               94.6 ms
 for_each(poolstl::par)/real_time                                   18.7 ms        0.044 ms           36
 for_each(std::execution::par)/real_time                            15.3 ms         12.9 ms           46
 sort()/real_time                                                    603 ms          602 ms            1
-sort(poolstl::par)/real_time                                        137 ms         11.8 ms            5
+sort(poolstl::par)/real_time                                        112 ms         6.64 ms            6
 sort(std::execution::par)/real_time                                 113 ms          102 ms            6
-pluggable_sort(poolstl::par, ..., pdqsort)/real_time               91.8 ms         11.9 ms            7
+pluggable_sort(poolstl::par, ..., pdqsort)/real_time               71.7 ms         6.67 ms           10
 transform()/real_time                                              95.0 ms         94.9 ms            7
 transform(poolstl::par)/real_time                                  17.4 ms        0.037 ms           38
 transform(std::execution::par)/real_time                           15.3 ms         13.2 ms           45

benchmark/algorithm_bench.cpp

@@ -93,6 +93,35 @@ BENCHMARK(for_each<std_par>)->Name("for_each(std::execution::par)")->UseRealTime
 
 ////////////////////////////////
 
+template <class ExecPolicy>
+void partition(benchmark::State& state) {
+    auto values = iota_vector<int>(arr_length);
+    std::vector<int> haystack(arr_length);
+
+    double pivot_frac = ((double)state.range(0)) / 100;
+    int pivot_val = (int)((double)values.size() * pivot_frac);
+    auto pred = [pivot_val] (const int& em) { return em < pivot_val; };
+
+    for ([[maybe_unused]] auto _ : state) {
+        if constexpr (is_policy<ExecPolicy>::value) {
+            auto res = std::partition(policy<ExecPolicy>::get(), values.begin(), values.end(), pred);
+            benchmark::DoNotOptimize(res);
+        } else {
+            auto res = std::partition(values.begin(), values.end(), pred);
+            benchmark::DoNotOptimize(res);
+        }
+        benchmark::ClobberMemory();
+    }
+}
+
+BENCHMARK(partition<seq>)->Name("partition()")->UseRealTime()->ArgName("pivot_percentile")->Arg(50);
+BENCHMARK(partition<poolstl_par>)->Name("partition(poolstl::par)")->UseRealTime()->ArgName("pivot_percentile")->Arg(50);
+#ifdef POOLSTL_BENCH_STD_PAR
+BENCHMARK(partition<std_par>)->Name("partition(std::execution::par)")->UseRealTime()->ArgName("pivot_percentile")->Arg(50);
+#endif
+
+////////////////////////////////
+
 template <class ExecPolicy>
 void sort(benchmark::State& state) {
     auto source = random_vector<int>(arr_length / 10);

include/poolstl/algorithm

@@ -12,6 +12,66 @@
 #include "internal/ttp_impl.hpp"
 #include "internal/thread_impl.hpp"
 
+namespace poolstl {
+    /**
+     * NOTE: Iterators are expected to be random access.
+     *
+     * Like `std::sort`, but allows specifying the sequential sort method, which must have the
+     * same signature as the comparator version of `std::sort`.
+     *
+     * Implemented as a high-level quicksort that delegates to `sort_func`, in parallel, once the range has been
+     * sufficiently partitioned.
+     */
+    template <class ExecPolicy, class RandIt, class Compare>
+    poolstl::internal::enable_if_poolstl_policy<ExecPolicy, void>
+    pluggable_sort(ExecPolicy &&policy, RandIt first, RandIt last, Compare comp,
+                   void (sort_func)(RandIt, RandIt, Compare) = std::sort) {
+        if (poolstl::internal::is_seq<ExecPolicy>(policy)) {
+            sort_func(first, last, comp);
+            return;
+        }
+
+        // Parallel partition.
+        // The partition_p2 method spawns and waits for its own child task. A deadlock is possible if all worker
+        // threads are waiting for tasks that in turn have to workers to execute them. This is only an issue because
+        // our thread pool does not have the concept of dependencies.
+        // So ensure
+        auto& task_pool = *policy.pool();
+        std::atomic<int> allowed_parallel_partitions{(int)task_pool.get_num_threads() / 2};
+
+        auto part_func = [&task_pool, &allowed_parallel_partitions](RandIt chunk_first, RandIt chunk_last,
+                                   poolstl::internal::pivot_predicate<Compare,
+                                   typename std::iterator_traits<RandIt>::value_type> pred) {
+            if (allowed_parallel_partitions.fetch_sub(1) > 0) {
+                return poolstl::internal::partition_p2(task_pool, chunk_first, chunk_last, pred);
+            } else {
+                return std::partition(chunk_first, chunk_last, pred);
+            }
+        };
+
+        poolstl::internal::parallel_quicksort(std::forward<ExecPolicy>(policy), first, last, comp, sort_func, part_func,
+                                              poolstl::internal::quicksort_pivot<RandIt>);
+    }
+
+    /**
+     * NOTE: Iterators are expected to be random access.
+     *
+     * Like `std::sort`, but allows specifying the sequential sort method, which must have the
+     * same signature as the comparator version of `std::sort`.
+     *
+     * Implemented as a parallel high-level quicksort that delegates to `sort_func` once the range has been
+     * sufficiently partitioned.
+     */
+    template <class ExecPolicy, class RandIt>
+    poolstl::internal::enable_if_poolstl_policy<ExecPolicy, void>
+    pluggable_sort(ExecPolicy &&policy, RandIt first, RandIt last,
+                   void (sort_func)(RandIt, RandIt,
+                                    std::less<typename std::iterator_traits<RandIt>::value_type>) = std::sort){
+        using T = typename std::iterator_traits<RandIt>::value_type;
+        pluggable_sort(std::forward<ExecPolicy>(policy), first, last, std::less<T>(), sort_func);
+    }
+}
+
 namespace std {
 
     /**
@@ -209,6 +269,22 @@ namespace std {
         return last;
     }
 
+    /**
+     * NOTE: Iterators are expected to be random access.
+     * See std::partition https://en.cppreference.com/w/cpp/algorithm/partition
+     *
+     * Current implementation uses at most 2 threads.
+     */
+    template <class ExecPolicy, class RandIt, class Predicate>
+    poolstl::internal::enable_if_poolstl_policy<ExecPolicy, RandIt>
+    partition(ExecPolicy &&policy, RandIt first, RandIt last, Predicate pred) {
+        if (poolstl::internal::is_seq<ExecPolicy>(policy)) {
+            return std::partition(first, last, pred);
+        }
+
+        return poolstl::internal::partition_p2(*policy.pool(), first, last, pred);
+    }
+
     /**
      * NOTE: Iterators are expected to be random access.
      * See std::sort https://en.cppreference.com/w/cpp/algorithm/sort
@@ -221,11 +297,7 @@ namespace std {
             return;
         }
 
-        poolstl::internal::parallel_quicksort(std::forward<ExecPolicy>(policy), first, last, comp,
-                                              std::sort<RandIt, Compare>,
-                                              std::partition<RandIt, poolstl::internal::pivot_predicate<Compare,
-                                                  typename std::iterator_traits<RandIt>::value_type>>,
-                                              poolstl::internal::quicksort_pivot<RandIt>);
+        poolstl::pluggable_sort(std::forward<ExecPolicy>(policy), first, last, comp, std::sort<RandIt, Compare>);
     }
 
     /**
@@ -377,48 +449,6 @@ namespace poolstl {
                                                      (void*)nullptr, 1, construct, f);
     }
 
-    /**
-     * NOTE: Iterators are expected to be random access.
-     *
-     * Like `std::sort`, but allows specifying the sequential sort method, which must have the
-     * same signature as the comparator version of `std::sort`.
-     *
-     * Implemented as a high-level quicksort that delegates to `sort_func`, in parallel, once the range has been
-     * sufficiently partitioned.
-     */
-    template <class ExecPolicy, class RandIt, class Compare>
-    poolstl::internal::enable_if_poolstl_policy<ExecPolicy, void>
-    pluggable_sort(ExecPolicy &&policy, RandIt first, RandIt last, Compare comp,
-                   void (sort_func)(RandIt, RandIt, Compare) = std::sort) {
-        if (poolstl::internal::is_seq<ExecPolicy>(policy)) {
-            sort_func(first, last, comp);
-            return;
-        }
-
-        poolstl::internal::parallel_quicksort(std::forward<ExecPolicy>(policy), first, last, comp, sort_func,
-                                              std::partition<RandIt, poolstl::internal::pivot_predicate<Compare,
-                                                  typename std::iterator_traits<RandIt>::value_type>>,
-                                              poolstl::internal::quicksort_pivot<RandIt>);
-    }
-
-    /**
-     * NOTE: Iterators are expected to be random access.
-     *
-     * Like `std::sort`, but allows specifying the sequential sort method, which must have the
-     * same signature as the comparator version of `std::sort`.
-     *
-     * Implemented as a parallel high-level quicksort that delegates to `sort_func` once the range has been
-     * sufficiently partitioned.
-     */
-    template <class ExecPolicy, class RandIt>
-    poolstl::internal::enable_if_poolstl_policy<ExecPolicy, void>
-    pluggable_sort(ExecPolicy &&policy, RandIt first, RandIt last,
-                   void (sort_func)(RandIt, RandIt,
-                                    std::less<typename std::iterator_traits<RandIt>::value_type>) = std::sort){
-        using T = typename std::iterator_traits<RandIt>::value_type;
-        pluggable_sort(std::forward<ExecPolicy>(policy), first, last, std::less<T>(), sort_func);
-    }
-
     /**
      * NOTE: Iterators are expected to be random access.
      *

include/poolstl/internal/ttp_impl.hpp

@@ -273,9 +273,14 @@ namespace poolstl {
             // Target partition size. Range will be recursively partitioned into partitions no bigger than this
             // size. Target approximately twice as many partitions as threads to reduce impact of uneven pivot
             // selection.
-            std::ptrdiff_t target_leaf_size = std::max(std::distance(first, last) / (task_pool.get_num_threads() * 2),
+            auto num_threads = task_pool.get_num_threads();
+            std::ptrdiff_t target_leaf_size = std::max(std::distance(first, last) / (num_threads * 2),
                                                        (std::ptrdiff_t)5);
 
+            if (num_threads == 1) {
+                target_leaf_size = std::distance(first, last);
+            }
+
             // task_thread_pool does not support creating task DAGs, so organize the code such that
             // all parallel tasks are independent. The parallel tasks can spawn additional parallel tasks, and they
             // record their "child" task's std::future into a common vector to be waited on by the main thread.
@@ -304,6 +309,39 @@ namespace poolstl {
             // Wait on all the parallel tasks.
             get_futures(futures);
         }
+
+        /**
+         * Partition range according to predicate. Unstable.
+         *
+         * This implementation only parallelizes with p=2; will spawn and wait for only one task.
+         */
+        template <class RandIt, class Predicate>
+        RandIt partition_p2(task_thread_pool::task_thread_pool &task_pool, RandIt first, RandIt last, Predicate pred) {
+            auto range_size = std::distance(first, last);
+            if (range_size < 4) {
+                return std::partition(first, last, pred);
+            }
+
+            // approach should be generalizable to arbitrary p
+
+            RandIt mid = std::next(first + range_size / 2);
+
+            // partition left and right halves in parallel
+            auto left_future = task_pool.submit(std::partition<RandIt, Predicate>, first, mid, pred);
+            RandIt right_mid = std::partition(mid, last, pred);
+            RandIt left_mid = left_future.get();
+
+            // merge the two partitioned halves
+            auto left_highs_size = std::distance(left_mid, mid);
+            auto right_lows_size = std::distance(mid, right_mid);
+            if (left_highs_size <= right_lows_size) {
+                std::swap_ranges(left_mid, mid, right_mid - left_highs_size);
+                return right_mid - left_highs_size;
+            } else {
+                std::swap_ranges(mid, right_mid, left_mid);
+                return left_mid + right_lows_size;
+            }
+        }
     }
 }
 

include/poolstl/seq_fwd.hpp

@@ -97,7 +97,9 @@ namespace std {
     POOLSTL_DEFINE_SEQ_FWD(std, for_each_n)
 #endif
 
+    POOLSTL_DEFINE_SEQ_FWD(std, partition)
     POOLSTL_DEFINE_SEQ_FWD(std, transform)
+    POOLSTL_DEFINE_SEQ_FWD(std, sort)
 
     // <numeric>
 

include/poolstl/variant_policy.hpp

@@ -105,7 +105,9 @@ namespace std {
     POOLSTL_DEFINE_PAR_IF_FWD(std, for_each_n)
 #endif
 
+    POOLSTL_DEFINE_PAR_IF_FWD(std, partition)
     POOLSTL_DEFINE_PAR_IF_FWD(std, transform)
+    POOLSTL_DEFINE_PAR_IF_FWD(std, sort)
 
     // <numeric>
 
@@ -121,6 +123,7 @@ namespace poolstl {
     // <poolstl/algorithm>
 
     POOLSTL_DEFINE_PAR_IF_FWD_VOID(poolstl, for_each_chunk)
+    POOLSTL_DEFINE_PAR_IF_FWD_VOID(poolstl, pluggable_sort)
 }
 #endif
 

tests/poolstl_test.cpp

@@ -321,16 +321,70 @@ TEST_CASE("inplace_merge", "[alg][algorithm]") {
     }
 }
 
+TEST_CASE("partition", "[alg][algorithm]") {
+    for (auto num_threads: test_thread_counts) {
+        ttp::task_thread_pool pool(num_threads);
+
+        for (auto num_iters: test_arr_sizes) {
+            for (int scramble_type = 0; scramble_type <= 2; ++scramble_type) {
+                auto source = iota_vector(num_iters);
+                switch (scramble_type) {
+                    case 0:
+                        std::reverse(source.begin(), source.end());
+                        break;
+                    case 1:
+                        scramble(source);
+                        break;
+                    default:
+                        break;
+                }
+
+                std::vector<int> pivots = {0, -1};
+                if (source.size() > 1) {
+                    pivots.push_back(*std::prev(source.end()));
+                    pivots.push_back(source[source.size() / 2]);
+                }
+
+                for (auto pivot : pivots) {
+                    auto pred = [pivot] (const int& em) { return em < pivot; };
+                    std::ptrdiff_t expected_left_size;
+                    {
+                        std::vector<int> work(source);
+                        auto ret = std::partition(work.begin(), work.end(), pred);
+                        expected_left_size = std::distance(work.begin(), ret);
+                    }
+                    {
+                        std::vector<int> work(source);
+                        auto mid = std::partition(poolstl::par_if(false), work.begin(), work.end(), pred);
+                        REQUIRE(expected_left_size == std::distance(work.begin(), mid));
+                        REQUIRE(std::is_partitioned(work.begin(), work.end(), pred));
+                    }
+                    {
+                        std::vector<int> work(source);
+                        auto mid = std::partition(poolstl::par.on(pool), work.begin(), work.end(), pred);
+                        REQUIRE(expected_left_size == std::distance(work.begin(), mid));
+                        REQUIRE(std::is_partitioned(work.begin(), work.end(), pred));
+                    }
+                }
+            }
+        }
+    }
+}
+
 TEST_CASE("sort", "[alg][algorithm]") {
     for (auto num_threads : test_thread_counts) {
         ttp::task_thread_pool pool(num_threads);
 
         for (auto num_iters : test_arr_sizes) {
-            for (int scramble_type = 0; scramble_type <= 2; ++scramble_type) {
+            for (int scramble_type = 0; scramble_type <= 3; ++scramble_type) {
                 auto source = iota_vector(num_iters);
                 switch (scramble_type) {
                     case 0: std::reverse(source.begin(), source.end()); break;
                     case 1: scramble(source); break;
+                    case 2:
+                        if (source.size() > 2) {
+                            std::swap(source[0], source[1]);
+                        }
                     default: break;
                 }