xref: /llvm-project/libcxx/include/__pstl/cpu_algos/transform.h (revision bbff52bfd49336bc0fdc83d8dfc616266bc07cbf)
1 //===----------------------------------------------------------------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #ifndef _LIBCPP___PSTL_CPU_ALGOS_TRANSFORM_H
10 #define _LIBCPP___PSTL_CPU_ALGOS_TRANSFORM_H
11 
12 #include <__algorithm/transform.h>
13 #include <__assert>
14 #include <__config>
15 #include <__iterator/concepts.h>
16 #include <__iterator/iterator_traits.h>
17 #include <__pstl/backend_fwd.h>
18 #include <__pstl/cpu_algos/cpu_traits.h>
19 #include <__type_traits/is_execution_policy.h>
20 #include <__utility/move.h>
21 #include <optional>
22 
23 #if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)
24 #  pragma GCC system_header
25 #endif
26 
27 _LIBCPP_PUSH_MACROS
28 #include <__undef_macros>
29 
30 #if _LIBCPP_STD_VER >= 17
31 
32 _LIBCPP_BEGIN_NAMESPACE_STD
33 namespace __pstl {
34 
35 template <class _Iterator1, class _DifferenceType, class _Iterator2, class _Function>
36 _LIBCPP_HIDE_FROM_ABI _Iterator2
37 __simd_transform(_Iterator1 __first1, _DifferenceType __n, _Iterator2 __first2, _Function __f) noexcept {
38   _PSTL_PRAGMA_SIMD
39   for (_DifferenceType __i = 0; __i < __n; ++__i)
40     __f(__first1[__i], __first2[__i]);
41   return __first2 + __n;
42 }
43 
44 template <class _Iterator1, class _DifferenceType, class _Iterator2, class _Iterator3, class _Function>
45 _LIBCPP_HIDE_FROM_ABI _Iterator3 __simd_transform(
46     _Iterator1 __first1, _DifferenceType __n, _Iterator2 __first2, _Iterator3 __first3, _Function __f) noexcept {
47   _PSTL_PRAGMA_SIMD
48   for (_DifferenceType __i = 0; __i < __n; ++__i)
49     __f(__first1[__i], __first2[__i], __first3[__i]);
50   return __first3 + __n;
51 }
52 
53 template <class _Backend, class _RawExecutionPolicy>
54 struct __cpu_parallel_transform {
55   template <class _Policy, class _ForwardIterator, class _ForwardOutIterator, class _UnaryOperation>
56   _LIBCPP_HIDE_FROM_ABI optional<_ForwardOutIterator>
57   operator()(_Policy&& __policy,
58              _ForwardIterator __first,
59              _ForwardIterator __last,
60              _ForwardOutIterator __result,
61              _UnaryOperation __op) const noexcept {
62     if constexpr (__is_parallel_execution_policy_v<_RawExecutionPolicy> &&
63                   __has_random_access_iterator_category_or_concept<_ForwardIterator>::value &&
64                   __has_random_access_iterator_category_or_concept<_ForwardOutIterator>::value) {
65       __cpu_traits<_Backend>::__for_each(
66           __first,
67           __last,
68           [&__policy, __op, __first, __result](_ForwardIterator __brick_first, _ForwardIterator __brick_last) {
69             using _TransformUnseq = __pstl::__transform<_Backend, __remove_parallel_policy_t<_RawExecutionPolicy>>;
70             auto __res            = _TransformUnseq()(
71                 std::__remove_parallel_policy(__policy),
72                 __brick_first,
73                 __brick_last,
74                 __result + (__brick_first - __first),
75                 __op);
76             _LIBCPP_ASSERT_INTERNAL(__res, "unseq/seq should never try to allocate!");
77             return *std::move(__res);
78           });
79       return __result + (__last - __first);
80     } else if constexpr (__is_unsequenced_execution_policy_v<_RawExecutionPolicy> &&
81                          __has_random_access_iterator_category_or_concept<_ForwardIterator>::value &&
82                          __has_random_access_iterator_category_or_concept<_ForwardOutIterator>::value) {
83       return __pstl::__simd_transform(
84           __first,
85           __last - __first,
86           __result,
87           [&](__iter_reference<_ForwardIterator> __in_value, __iter_reference<_ForwardOutIterator> __out_value) {
88             __out_value = __op(__in_value);
89           });
90     } else {
91       return std::transform(__first, __last, __result, __op);
92     }
93   }
94 };
95 
96 template <class _Backend, class _RawExecutionPolicy>
97 struct __cpu_parallel_transform_binary {
98   template <class _Policy,
99             class _ForwardIterator1,
100             class _ForwardIterator2,
101             class _ForwardOutIterator,
102             class _BinaryOperation>
103   _LIBCPP_HIDE_FROM_ABI optional<_ForwardOutIterator>
104   operator()(_Policy&& __policy,
105              _ForwardIterator1 __first1,
106              _ForwardIterator1 __last1,
107              _ForwardIterator2 __first2,
108              _ForwardOutIterator __result,
109              _BinaryOperation __op) const noexcept {
110     if constexpr (__is_parallel_execution_policy_v<_RawExecutionPolicy> &&
111                   __has_random_access_iterator_category_or_concept<_ForwardIterator1>::value &&
112                   __has_random_access_iterator_category_or_concept<_ForwardIterator2>::value &&
113                   __has_random_access_iterator_category_or_concept<_ForwardOutIterator>::value) {
114       auto __res = __cpu_traits<_Backend>::__for_each(
115           __first1,
116           __last1,
117           [&__policy, __op, __first1, __first2, __result](
118               _ForwardIterator1 __brick_first, _ForwardIterator1 __brick_last) {
119             using _TransformBinaryUnseq =
120                 __pstl::__transform_binary<_Backend, __remove_parallel_policy_t<_RawExecutionPolicy>>;
121             return _TransformBinaryUnseq()(
122                 std::__remove_parallel_policy(__policy),
123                 __brick_first,
124                 __brick_last,
125                 __first2 + (__brick_first - __first1),
126                 __result + (__brick_first - __first1),
127                 __op);
128           });
129       if (!__res)
130         return nullopt;
131       return __result + (__last1 - __first1);
132     } else if constexpr (__is_unsequenced_execution_policy_v<_RawExecutionPolicy> &&
133                          __has_random_access_iterator_category_or_concept<_ForwardIterator1>::value &&
134                          __has_random_access_iterator_category_or_concept<_ForwardIterator2>::value &&
135                          __has_random_access_iterator_category_or_concept<_ForwardOutIterator>::value) {
136       return __pstl::__simd_transform(
137           __first1,
138           __last1 - __first1,
139           __first2,
140           __result,
141           [&](__iter_reference<_ForwardIterator1> __in1,
142               __iter_reference<_ForwardIterator2> __in2,
143               __iter_reference<_ForwardOutIterator> __out_value) { __out_value = __op(__in1, __in2); });
144     } else {
145       return std::transform(__first1, __last1, __first2, __result, __op);
146     }
147   }
148 };
149 
150 } // namespace __pstl
151 _LIBCPP_END_NAMESPACE_STD
152 
153 #endif // _LIBCPP_STD_VER >= 17
154 
155 _LIBCPP_POP_MACROS
156 
157 #endif // _LIBCPP___PSTL_CPU_ALGOS_TRANSFORM_H
158