[libc++] Overhaul the PSTL dispatching mechanism (llvm#88131)

The experimental PSTL's current dispatching mechanism was designed with
flexibility in mind. However, while reviewing the in-progress OpenMP
backend, I realized that the dispatching mechanism based on ADL and
default definitions in the frontend had several downsides. To name a
few:

1. The dispatching of an algorithm to the back-end and its default
   implementation is bundled together via `_LIBCPP_PSTL_CUSTOMIZATION_POINT`.
   This makes the dispatching really confusing and leads to annoyances
   such as variable shadowing and weird lambda captures in the front-end.
2. The distinction between back-end functions and front-end algorithms
   is not as clear as it could be, which led us to call one where we meant
   the other in a few cases. This is bad due to the exception requirements
   of the PSTL: calling a front-end algorithm inside the implementation of
   a back-end is incorrect for exception-safety.
3. There are two levels of back-end dispatching in the PSTL, which treat
   CPU backends as a special case. This was confusing and not as flexible
   as we'd like. For example, there was no straightforward way to dispatch
   all uses of `unseq` to a specific back-end from the OpenMP backend,
   or for CPU backends to fall back on each other.

This patch rewrites the backend dispatching mechanism to solve these
problems, but doesn't touch any of the actual implementation of
algorithms. Specifically, this rewrite has the following
characteristics:

- There is a single level of backend dispatching, however partial backends can
  be stacked to provide a full implementation of the PSTL. The two-level dispatching
  that was used for CPU-based backends is handled by providing CPU-based basis 
  operations as simple helpers that can easily be reused when defining any PSTL 
  backend.

- The default definitions for algorithms are separated from their dispatching logic.

- The front-end is thus simplified a whole lot and made very consistent
  for all algorithms, which makes it easier to audit the front-end for
  things like exception-correctness, appropriate forwarding, etc.

Fixes llvm#70718

Author: ldionne

libcxx/benchmarks/algorithms/pstl.stable_sort.bench.cpp

@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include <algorithm>
+#include <execution>
 
 #include "common.h"
 

libcxx/include/CMakeLists.txt

@@ -73,7 +73,6 @@ set(files
   __algorithm/pop_heap.h
   __algorithm/prev_permutation.h
   __algorithm/pstl.h
-  __algorithm/pstl_frontend_dispatch.h
   __algorithm/push_heap.h
   __algorithm/ranges_adjacent_find.h
   __algorithm/ranges_all_of.h
@@ -576,11 +575,12 @@ set(files
   __numeric/transform_reduce.h
   __ostream/basic_ostream.h
   __ostream/print.h
+  __pstl/backend.h
+  __pstl/backend_fwd.h
+  __pstl/backends/default.h
   __pstl/backends/libdispatch.h
   __pstl/backends/serial.h
   __pstl/backends/std_thread.h
-  __pstl/configuration.h
-  __pstl/configuration_fwd.h
   __pstl/cpu_algos/any_of.h
   __pstl/cpu_algos/cpu_traits.h
   __pstl/cpu_algos/fill.h
@@ -590,6 +590,8 @@ set(files
   __pstl/cpu_algos/stable_sort.h
   __pstl/cpu_algos/transform.h
   __pstl/cpu_algos/transform_reduce.h
+  __pstl/dispatch.h
+  __pstl/handle_exception.h
   __random/bernoulli_distribution.h
   __random/binomial_distribution.h
   __random/cauchy_distribution.h

libcxx/include/__algorithm/pstl.h

@@ -9,17 +9,19 @@
 #ifndef _LIBCPP___NUMERIC_PSTL_H
 #define _LIBCPP___NUMERIC_PSTL_H
 
-#include <__algorithm/pstl_frontend_dispatch.h>
 #include <__config>
 #include <__functional/identity.h>
 #include <__functional/operations.h>
 #include <__iterator/cpp17_iterator_concepts.h>
 #include <__iterator/iterator_traits.h>
-#include <__numeric/transform_reduce.h>
-#include <__pstl/configuration.h>
+#include <__pstl/backend.h>
+#include <__pstl/dispatch.h>
+#include <__pstl/handle_exception.h>
+#include <__type_traits/enable_if.h>
 #include <__type_traits/is_execution_policy.h>
+#include <__type_traits/remove_cvref.h>
+#include <__utility/forward.h>
 #include <__utility/move.h>
-#include <optional>
 
 #if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)
 #  pragma GCC system_header
@@ -33,30 +35,50 @@ _LIBCPP_PUSH_MACROS
 _LIBCPP_BEGIN_NAMESPACE_STD
 
 template <class _ExecutionPolicy,
-          class _ForwardIterator1,
-          class _ForwardIterator2,
+          class _ForwardIterator,
           class _Tp,
-          class _BinaryOperation1,
-          class _BinaryOperation2,
+          class _BinaryOperation,
           class _RawPolicy                                    = __remove_cvref_t<_ExecutionPolicy>,
           enable_if_t<is_execution_policy_v<_RawPolicy>, int> = 0>
-_LIBCPP_HIDE_FROM_ABI optional<_Tp> __transform_reduce(
-    _ExecutionPolicy&&,
-    _ForwardIterator1&& __first1,
-    _ForwardIterator1&& __last1,
-    _ForwardIterator2&& __first2,
-    _Tp&& __init,
-    _BinaryOperation1&& __reduce,
-    _BinaryOperation2&& __transform) noexcept {
-  using _Backend = typename __select_backend<_RawPolicy>::type;
-  return std::__pstl_transform_reduce<_RawPolicy>(
-      _Backend{},
-      std::move(__first1),
-      std::move(__last1),
-      std::move(__first2),
+_LIBCPP_HIDE_FROM_ABI _Tp reduce(
+    _ExecutionPolicy&& __policy, _ForwardIterator __first, _ForwardIterator __last, _Tp __init, _BinaryOperation __op) {
+  _LIBCPP_REQUIRE_CPP17_FORWARD_ITERATOR(_ForwardIterator, "reduce requires ForwardIterators");
+  using _Implementation = __pstl::__dispatch<__pstl::__reduce, __pstl::__current_configuration, _RawPolicy>;
+  return __pstl::__handle_exception<_Implementation>(
+      std::forward<_ExecutionPolicy>(__policy),
+      std::move(__first),
+      std::move(__last),
       std::move(__init),
-      std::move(__reduce),
-      std::move(__transform));
+      std::move(__op));
+}
+
+template <class _ExecutionPolicy,
+          class _ForwardIterator,
+          class _Tp,
+          class _RawPolicy                                    = __remove_cvref_t<_ExecutionPolicy>,
+          enable_if_t<is_execution_policy_v<_RawPolicy>, int> = 0>
+_LIBCPP_HIDE_FROM_ABI _Tp
+reduce(_ExecutionPolicy&& __policy, _ForwardIterator __first, _ForwardIterator __last, _Tp __init) {
+  _LIBCPP_REQUIRE_CPP17_FORWARD_ITERATOR(_ForwardIterator, "reduce requires ForwardIterators");
+  using _Implementation = __pstl::__dispatch<__pstl::__reduce, __pstl::__current_configuration, _RawPolicy>;
+  return __pstl::__handle_exception<_Implementation>(
+      std::forward<_ExecutionPolicy>(__policy), std::move(__first), std::move(__last), std::move(__init), plus{});
+}
+
+template <class _ExecutionPolicy,
+          class _ForwardIterator,
+          class _RawPolicy                                    = __remove_cvref_t<_ExecutionPolicy>,
+          enable_if_t<is_execution_policy_v<_RawPolicy>, int> = 0>
+_LIBCPP_HIDE_FROM_ABI __iter_value_type<_ForwardIterator>
+reduce(_ExecutionPolicy&& __policy, _ForwardIterator __first, _ForwardIterator __last) {
+  _LIBCPP_REQUIRE_CPP17_FORWARD_ITERATOR(_ForwardIterator, "reduce requires ForwardIterators");
+  using _Implementation = __pstl::__dispatch<__pstl::__reduce, __pstl::__current_configuration, _RawPolicy>;
+  return __pstl::__handle_exception<_Implementation>(
+      std::forward<_ExecutionPolicy>(__policy),
+      std::move(__first),
+      std::move(__last),
+      __iter_value_type<_ForwardIterator>(),
+      plus{});
 }
 
 template <class _ExecutionPolicy,
@@ -77,18 +99,16 @@ _LIBCPP_HIDE_FROM_ABI _Tp transform_reduce(
     _BinaryOperation2 __transform) {
   _LIBCPP_REQUIRE_CPP17_FORWARD_ITERATOR(_ForwardIterator1, "transform_reduce requires ForwardIterators");
   _LIBCPP_REQUIRE_CPP17_FORWARD_ITERATOR(_ForwardIterator2, "transform_reduce requires ForwardIterators");
-  auto __res = std::__transform_reduce(
-      __policy,
+  using _Implementation =
+      __pstl::__dispatch<__pstl::__transform_reduce_binary, __pstl::__current_configuration, _RawPolicy>;
+  return __pstl::__handle_exception<_Implementation>(
+      std::forward<_ExecutionPolicy>(__policy),
       std::move(__first1),
       std::move(__last1),
       std::move(__first2),
       std::move(__init),
       std::move(__reduce),
       std::move(__transform));
-
-  if (!__res)
-    std::__throw_bad_alloc();
-  return *std::move(__res);
 }
 
 // This overload doesn't get a customization point because it's trivial to detect (through e.g.
@@ -97,7 +117,8 @@ template <class _ExecutionPolicy,
           class _ForwardIterator1,
           class _ForwardIterator2,
           class _Tp,
-          enable_if_t<is_execution_policy_v<__remove_cvref_t<_ExecutionPolicy>>, int> = 0>
+          class _RawPolicy                                    = __remove_cvref_t<_ExecutionPolicy>,
+          enable_if_t<is_execution_policy_v<_RawPolicy>, int> = 0>
 _LIBCPP_HIDE_FROM_ABI _Tp transform_reduce(
     _ExecutionPolicy&& __policy,
     _ForwardIterator1 __first1,
@@ -106,31 +127,16 @@ _LIBCPP_HIDE_FROM_ABI _Tp transform_reduce(
     _Tp __init) {
   _LIBCPP_REQUIRE_CPP17_FORWARD_ITERATOR(_ForwardIterator1, "transform_reduce requires ForwardIterators");
   _LIBCPP_REQUIRE_CPP17_FORWARD_ITERATOR(_ForwardIterator2, "transform_reduce requires ForwardIterators");
-  return std::transform_reduce(__policy, __first1, __last1, __first2, __init, plus{}, multiplies{});
-}
-
-template <class _ExecutionPolicy,
-          class _ForwardIterator,
-          class _Tp,
-          class _BinaryOperation,
-          class _UnaryOperation,
-          class _RawPolicy                                    = __remove_cvref_t<_ExecutionPolicy>,
-          enable_if_t<is_execution_policy_v<_RawPolicy>, int> = 0>
-[[nodiscard]] _LIBCPP_HIDE_FROM_ABI optional<__remove_cvref_t<_Tp>> __transform_reduce(
-    _ExecutionPolicy&&,
-    _ForwardIterator&& __first,
-    _ForwardIterator&& __last,
-    _Tp&& __init,
-    _BinaryOperation&& __reduce,
-    _UnaryOperation&& __transform) noexcept {
-  using _Backend = typename __select_backend<_RawPolicy>::type;
-  return std::__pstl_transform_reduce<_RawPolicy>(
-      _Backend{},
-      std::move(__first),
-      std::move(__last),
+  using _Implementation =
+      __pstl::__dispatch<__pstl::__transform_reduce_binary, __pstl::__current_configuration, _RawPolicy>;
+  return __pstl::__handle_exception<_Implementation>(
+      std::forward<_ExecutionPolicy>(__policy),
+      std::move(__first1),
+      std::move(__last1),
+      std::move(__first2),
       std::move(__init),
-      std::move(__reduce),
-      std::move(__transform));
+      plus{},
+      multiplies{});
 }
 
 template <class _ExecutionPolicy,
@@ -148,86 +154,14 @@ _LIBCPP_HIDE_FROM_ABI _Tp transform_reduce(
     _BinaryOperation __reduce,
     _UnaryOperation __transform) {
   _LIBCPP_REQUIRE_CPP17_FORWARD_ITERATOR(_ForwardIterator, "transform_reduce requires ForwardIterators");
-  auto __res = std::__transform_reduce(
-      __policy, std::move(__first), std::move(__last), std::move(__init), std::move(__reduce), std::move(__transform));
-  if (!__res)
-    std::__throw_bad_alloc();
-  return *std::move(__res);
-}
-
-template <class>
-void __pstl_reduce();
-
-template <class _ExecutionPolicy,
-          class _ForwardIterator,
-          class _Tp,
-          class _BinaryOperation                              = plus<>,
-          class _RawPolicy                                    = __remove_cvref_t<_ExecutionPolicy>,
-          enable_if_t<is_execution_policy_v<_RawPolicy>, int> = 0>
-[[nodiscard]] _LIBCPP_HIDE_FROM_ABI optional<_Tp>
-__reduce(_ExecutionPolicy&& __policy,
-         _ForwardIterator&& __first,
-         _ForwardIterator&& __last,
-         _Tp&& __init,
-         _BinaryOperation&& __op = {}) noexcept {
-  return std::__pstl_frontend_dispatch(
-      _LIBCPP_PSTL_CUSTOMIZATION_POINT(__pstl_reduce, _RawPolicy),
-      [&__policy](_ForwardIterator __g_first, _ForwardIterator __g_last, _Tp __g_init, _BinaryOperation __g_op) {
-        return std::__transform_reduce(
-            __policy, std::move(__g_first), std::move(__g_last), std::move(__g_init), std::move(__g_op), __identity{});
-      },
+  using _Implementation = __pstl::__dispatch<__pstl::__transform_reduce, __pstl::__current_configuration, _RawPolicy>;
+  return __pstl::__handle_exception<_Implementation>(
+      std::forward<_ExecutionPolicy>(__policy),
       std::move(__first),
       std::move(__last),
       std::move(__init),
-      std::move(__op));
-}
-
-template <class _ExecutionPolicy,
-          class _ForwardIterator,
-          class _Tp,
-          class _BinaryOperation                              = plus<>,
-          class _RawPolicy                                    = __remove_cvref_t<_ExecutionPolicy>,
-          enable_if_t<is_execution_policy_v<_RawPolicy>, int> = 0>
-_LIBCPP_HIDE_FROM_ABI _Tp
-reduce(_ExecutionPolicy&& __policy,
-       _ForwardIterator __first,
-       _ForwardIterator __last,
-       _Tp __init,
-       _BinaryOperation __op = {}) {
-  _LIBCPP_REQUIRE_CPP17_FORWARD_ITERATOR(_ForwardIterator, "reduce requires ForwardIterators");
-  auto __res = std::__reduce(__policy, std::move(__first), std::move(__last), std::move(__init), std::move(__op));
-  if (!__res)
-    std::__throw_bad_alloc();
-  return *std::move(__res);
-}
-
-template <class _ExecutionPolicy,
-          class _ForwardIterator,
-          class _RawPolicy                                    = __remove_cvref_t<_ExecutionPolicy>,
-          enable_if_t<is_execution_policy_v<_RawPolicy>, int> = 0>
-[[nodiscard]] _LIBCPP_HIDE_FROM_ABI optional<__iter_value_type<_ForwardIterator>>
-__reduce(_ExecutionPolicy&& __policy, _ForwardIterator&& __first, _ForwardIterator&& __last) noexcept {
-  return std::__pstl_frontend_dispatch(
-      _LIBCPP_PSTL_CUSTOMIZATION_POINT(__pstl_reduce, _RawPolicy),
-      [&__policy](_ForwardIterator __g_first, _ForwardIterator __g_last) {
-        return std::__reduce(
-            __policy, std::move(__g_first), std::move(__g_last), __iter_value_type<_ForwardIterator>());
-      },
-      std::move(__first),
-      std::move(__last));
-}
-
-template <class _ExecutionPolicy,
-          class _ForwardIterator,
-          class _RawPolicy                                    = __remove_cvref_t<_ExecutionPolicy>,
-          enable_if_t<is_execution_policy_v<_RawPolicy>, int> = 0>
-_LIBCPP_HIDE_FROM_ABI __iter_value_type<_ForwardIterator>
-reduce(_ExecutionPolicy&& __policy, _ForwardIterator __first, _ForwardIterator __last) {
-  _LIBCPP_REQUIRE_CPP17_FORWARD_ITERATOR(_ForwardIterator, "reduce requires ForwardIterators");
-  auto __res = std::__reduce(__policy, std::move(__first), std::move(__last));
-  if (!__res)
-    std::__throw_bad_alloc();
-  return *std::move(__res);
+      std::move(__reduce),
+      std::move(__transform));
 }
 
 _LIBCPP_END_NAMESPACE_STD

libcxx/include/__algorithm/pstl_frontend_dispatch.h

@@ -6,22 +6,30 @@
 //
 //===----------------------------------------------------------------------===//
 
-#ifndef _LIBCPP___PSTL_CONFIGURATION_H
-#define _LIBCPP___PSTL_CONFIGURATION_H
+#ifndef _LIBCPP___PSTL_BACKEND_H
+#define _LIBCPP___PSTL_BACKEND_H
 
 #include <__config>
-#include <__pstl/configuration_fwd.h>
+#include <__pstl/backend_fwd.h>
 
 #if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)
 #  pragma GCC system_header
 #endif
 
+_LIBCPP_PUSH_MACROS
+#include <__undef_macros>
+
 #if defined(_LIBCPP_PSTL_BACKEND_SERIAL)
+#  include <__pstl/backends/default.h>
 #  include <__pstl/backends/serial.h>
 #elif defined(_LIBCPP_PSTL_BACKEND_STD_THREAD)
+#  include <__pstl/backends/default.h>
 #  include <__pstl/backends/std_thread.h>
 #elif defined(_LIBCPP_PSTL_BACKEND_LIBDISPATCH)
+#  include <__pstl/backends/default.h>
 #  include <__pstl/backends/libdispatch.h>
 #endif
 
-#endif // _LIBCPP___PSTL_CONFIGURATION_H
+_LIBCPP_POP_MACROS
+
+#endif // _LIBCPP___PSTL_BACKEND_H