220 lines · c
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-exception6//7//===----------------------------------------------------------------------===//8 9#ifndef _LIBCPP___PSTL_CPU_ALGOS_TRANSFORM_REDUCE_H10#define _LIBCPP___PSTL_CPU_ALGOS_TRANSFORM_REDUCE_H11 12#include <__assert>13#include <__config>14#include <__iterator/concepts.h>15#include <__iterator/iterator_traits.h>16#include <__numeric/transform_reduce.h>17#include <__pstl/backend_fwd.h>18#include <__pstl/cpu_algos/cpu_traits.h>19#include <__type_traits/desugars_to.h>20#include <__type_traits/is_arithmetic.h>21#include <__type_traits/is_execution_policy.h>22#include <__utility/move.h>23#include <optional>24 25#if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)26# pragma GCC system_header27#endif28 29_LIBCPP_PUSH_MACROS30#include <__undef_macros>31 32#if _LIBCPP_STD_VER >= 1733 34_LIBCPP_BEGIN_NAMESPACE_STD35namespace __pstl {36 37template <typename _Backend,38 typename _DifferenceType,39 typename _Tp,40 typename _BinaryOperation,41 typename _UnaryOperation,42 typename _UnaryResult = invoke_result_t<_UnaryOperation, _DifferenceType>,43 __enable_if_t<__desugars_to_v<__plus_tag, _BinaryOperation, _Tp, _UnaryResult> && is_arithmetic_v<_Tp> &&44 is_arithmetic_v<_UnaryResult>,45 int> = 0>46_LIBCPP_HIDE_FROM_ABI _Tp47__simd_transform_reduce(_DifferenceType __n, _Tp __init, _BinaryOperation, _UnaryOperation __f) noexcept {48 _PSTL_PRAGMA_SIMD_REDUCTION(+ : __init)49 for (_DifferenceType __i = 0; __i < __n; ++__i)50 __init += __f(__i);51 return __init;52}53 54template <typename _Backend,55 typename _Size,56 typename _Tp,57 typename _BinaryOperation,58 typename _UnaryOperation,59 typename _UnaryResult = invoke_result_t<_UnaryOperation, _Size>,60 __enable_if_t<!(__desugars_to_v<__plus_tag, _BinaryOperation, _Tp, _UnaryResult> && is_arithmetic_v<_Tp> &&61 is_arithmetic_v<_UnaryResult>),62 int> = 0>63_LIBCPP_HIDE_FROM_ABI _Tp64__simd_transform_reduce(_Size __n, _Tp __init, _BinaryOperation __binary_op, _UnaryOperation __f) noexcept {65 constexpr size_t __lane_size = __cpu_traits<_Backend>::__lane_size;66 const _Size __block_size = __lane_size / sizeof(_Tp);67 if (__n > 2 * __block_size && __block_size > 1) {68 alignas(__lane_size) char __lane_buffer[__lane_size];69 _Tp* __lane = reinterpret_cast<_Tp*>(__lane_buffer);70 71 // initializer72 _PSTL_PRAGMA_SIMD73 for (_Size __i = 0; __i < __block_size; ++__i) {74 ::new (__lane + __i) _Tp(__binary_op(__f(__i), __f(__block_size + __i)));75 }76 // main loop77 _Size __i = 2 * __block_size;78 const _Size __last_iteration = __block_size * (__n / __block_size);79 for (; __i < __last_iteration; __i += __block_size) {80 _PSTL_PRAGMA_SIMD81 for (_Size __j = 0; __j < __block_size; ++__j) {82 __lane[__j] = __binary_op(std::move(__lane[__j]), __f(__i + __j));83 }84 }85 // remainder86 _PSTL_PRAGMA_SIMD87 for (_Size __j = 0; __j < __n - __last_iteration; ++__j) {88 __lane[__j] = __binary_op(std::move(__lane[__j]), __f(__last_iteration + __j));89 }90 // combiner91 for (_Size __j = 0; __j < __block_size; ++__j) {92 __init = __binary_op(std::move(__init), std::move(__lane[__j]));93 }94 // destroyer95 _PSTL_PRAGMA_SIMD96 for (_Size __j = 0; __j < __block_size; ++__j) {97 __lane[__j].~_Tp();98 }99 } else {100 for (_Size __i = 0; __i < __n; ++__i) {101 __init = __binary_op(std::move(__init), __f(__i));102 }103 }104 return __init;105}106 107template <class _Backend, class _RawExecutionPolicy>108struct __cpu_parallel_transform_reduce_binary {109 template <class _Policy,110 class _ForwardIterator1,111 class _ForwardIterator2,112 class _Tp,113 class _BinaryOperation1,114 class _BinaryOperation2>115 _LIBCPP_HIDE_FROM_ABI optional<_Tp> operator()(116 _Policy&& __policy,117 _ForwardIterator1 __first1,118 _ForwardIterator1 __last1,119 _ForwardIterator2 __first2,120 _Tp __init,121 _BinaryOperation1 __reduce,122 _BinaryOperation2 __transform) const noexcept {123 if constexpr (__is_parallel_execution_policy_v<_RawExecutionPolicy> &&124 __has_random_access_iterator_category_or_concept<_ForwardIterator1>::value &&125 __has_random_access_iterator_category_or_concept<_ForwardIterator2>::value) {126 return __cpu_traits<_Backend>::__transform_reduce(127 __first1,128 std::move(__last1),129 [__first1, __first2, __transform](_ForwardIterator1 __iter) {130 return __transform(*__iter, *(__first2 + (__iter - __first1)));131 },132 std::move(__init),133 std::move(__reduce),134 [&__policy, __first1, __first2, __reduce, __transform](135 _ForwardIterator1 __brick_first, _ForwardIterator1 __brick_last, _Tp __brick_init) {136 using _TransformReduceBinaryUnseq =137 __pstl::__transform_reduce_binary<_Backend, __remove_parallel_policy_t<_RawExecutionPolicy>>;138 return *_TransformReduceBinaryUnseq()(139 std::__remove_parallel_policy(__policy),140 __brick_first,141 std::move(__brick_last),142 __first2 + (__brick_first - __first1),143 std::move(__brick_init),144 std::move(__reduce),145 std::move(__transform));146 });147 } else if constexpr (__is_unsequenced_execution_policy_v<_RawExecutionPolicy> &&148 __has_random_access_iterator_category_or_concept<_ForwardIterator1>::value &&149 __has_random_access_iterator_category_or_concept<_ForwardIterator2>::value) {150 return __pstl::__simd_transform_reduce<_Backend>(151 __last1 - __first1,152 std::move(__init),153 std::move(__reduce),154 [&](__iterator_difference_type<_ForwardIterator1> __i) { return __transform(__first1[__i], __first2[__i]); });155 } else {156 return std::transform_reduce(157 std::move(__first1),158 std::move(__last1),159 std::move(__first2),160 std::move(__init),161 std::move(__reduce),162 std::move(__transform));163 }164 }165};166 167template <class _Backend, class _RawExecutionPolicy>168struct __cpu_parallel_transform_reduce {169 template <class _Policy, class _ForwardIterator, class _Tp, class _BinaryOperation, class _UnaryOperation>170 _LIBCPP_HIDE_FROM_ABI optional<_Tp>171 operator()(_Policy&& __policy,172 _ForwardIterator __first,173 _ForwardIterator __last,174 _Tp __init,175 _BinaryOperation __reduce,176 _UnaryOperation __transform) const noexcept {177 if constexpr (__is_parallel_execution_policy_v<_RawExecutionPolicy> &&178 __has_random_access_iterator_category_or_concept<_ForwardIterator>::value) {179 return __cpu_traits<_Backend>::__transform_reduce(180 std::move(__first),181 std::move(__last),182 [__transform](_ForwardIterator __iter) { return __transform(*__iter); },183 std::move(__init),184 __reduce,185 [&__policy, __transform, __reduce](auto __brick_first, auto __brick_last, _Tp __brick_init) {186 using _TransformReduceUnseq =187 __pstl::__transform_reduce<_Backend, __remove_parallel_policy_t<_RawExecutionPolicy>>;188 auto __res = _TransformReduceUnseq()(189 std::__remove_parallel_policy(__policy),190 std::move(__brick_first),191 std::move(__brick_last),192 std::move(__brick_init),193 std::move(__reduce),194 std::move(__transform));195 _LIBCPP_ASSERT_INTERNAL(__res, "unseq/seq should never try to allocate!");196 return *std::move(__res);197 });198 } else if constexpr (__is_unsequenced_execution_policy_v<_RawExecutionPolicy> &&199 __has_random_access_iterator_category_or_concept<_ForwardIterator>::value) {200 return __pstl::__simd_transform_reduce<_Backend>(201 __last - __first,202 std::move(__init),203 std::move(__reduce),204 [=, &__transform](__iterator_difference_type<_ForwardIterator> __i) { return __transform(__first[__i]); });205 } else {206 return std::transform_reduce(207 std::move(__first), std::move(__last), std::move(__init), std::move(__reduce), std::move(__transform));208 }209 }210};211 212} // namespace __pstl213_LIBCPP_END_NAMESPACE_STD214 215#endif // _LIBCPP_STD_VER >= 17216 217_LIBCPP_POP_MACROS218 219#endif // _LIBCPP___PSTL_CPU_ALGOS_TRANSFORM_REDUCE_H220