__pstl/cpu_algos/transform_reduce.h

0001 //===----------------------------------------------------------------------===//
0002 //
0003 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
0004 // See https://llvm.org/LICENSE.txt for license information.
0005 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
0006 //
0007 //===----------------------------------------------------------------------===//
0008
0009 #ifndef _LIBCPP___PSTL_CPU_ALGOS_TRANSFORM_REDUCE_H
0010 #define _LIBCPP___PSTL_CPU_ALGOS_TRANSFORM_REDUCE_H
0011
0012 #include <__assert>
0013 #include <__config>
0014 #include <__iterator/concepts.h>
0015 #include <__iterator/iterator_traits.h>
0016 #include <__numeric/transform_reduce.h>
0017 #include <__pstl/backend_fwd.h>
0018 #include <__pstl/cpu_algos/cpu_traits.h>
0019 #include <__type_traits/desugars_to.h>
0020 #include <__type_traits/is_arithmetic.h>
0021 #include <__type_traits/is_execution_policy.h>
0022 #include <__utility/move.h>
0023 #include <optional>
0024
0025 #if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)
0026 #  pragma GCC system_header
0027 #endif
0028
0029 _LIBCPP_PUSH_MACROS
0030 #include <__undef_macros>
0031
0032 #if _LIBCPP_STD_VER >= 17
0033
0034 _LIBCPP_BEGIN_NAMESPACE_STD
0035 namespace __pstl {
0036
0037 template <typename _Backend,
0038           typename _DifferenceType,
0039           typename _Tp,
0040           typename _BinaryOperation,
0041           typename _UnaryOperation,
0042           typename _UnaryResult = invoke_result_t<_UnaryOperation, _DifferenceType>,
0043           __enable_if_t<__desugars_to_v<__plus_tag, _BinaryOperation, _Tp, _UnaryResult> && is_arithmetic_v<_Tp> &&
0044                             is_arithmetic_v<_UnaryResult>,
0045                         int>    = 0>
0046 _LIBCPP_HIDE_FROM_ABI _Tp
0047 __simd_transform_reduce(_DifferenceType __n, _Tp __init, _BinaryOperation, _UnaryOperation __f) noexcept {
0048   _PSTL_PRAGMA_SIMD_REDUCTION(+ : __init)
0049   for (_DifferenceType __i = 0; __i < __n; ++__i)
0050     __init += __f(__i);
0051   return __init;
0052 }
0053
0054 template <typename _Backend,
0055           typename _Size,
0056           typename _Tp,
0057           typename _BinaryOperation,
0058           typename _UnaryOperation,
0059           typename _UnaryResult = invoke_result_t<_UnaryOperation, _Size>,
0060           __enable_if_t<!(__desugars_to_v<__plus_tag, _BinaryOperation, _Tp, _UnaryResult> && is_arithmetic_v<_Tp> &&
0061                           is_arithmetic_v<_UnaryResult>),
0062                         int>    = 0>
0063 _LIBCPP_HIDE_FROM_ABI _Tp
0064 __simd_transform_reduce(_Size __n, _Tp __init, _BinaryOperation __binary_op, _UnaryOperation __f) noexcept {
0065   constexpr size_t __lane_size = __cpu_traits<_Backend>::__lane_size;
0066   const _Size __block_size     = __lane_size / sizeof(_Tp);
0067   if (__n > 2 * __block_size && __block_size > 1) {
0068     alignas(__lane_size) char __lane_buffer[__lane_size];
0069     _Tp* __lane = reinterpret_cast<_Tp*>(__lane_buffer);
0070
0071     // initializer
0072     _PSTL_PRAGMA_SIMD
0073     for (_Size __i = 0; __i < __block_size; ++__i) {
0074       ::new (__lane + __i) _Tp(__binary_op(__f(__i), __f(__block_size + __i)));
0075     }
0076     // main loop
0077     _Size __i                    = 2 * __block_size;
0078     const _Size __last_iteration = __block_size * (__n / __block_size);
0079     for (; __i < __last_iteration; __i += __block_size) {
0080       _PSTL_PRAGMA_SIMD
0081       for (_Size __j = 0; __j < __block_size; ++__j) {
0082         __lane[__j] = __binary_op(std::move(__lane[__j]), __f(__i + __j));
0083       }
0084     }
0085     // remainder
0086     _PSTL_PRAGMA_SIMD
0087     for (_Size __j = 0; __j < __n - __last_iteration; ++__j) {
0088       __lane[__j] = __binary_op(std::move(__lane[__j]), __f(__last_iteration + __j));
0089     }
0090     // combiner
0091     for (_Size __j = 0; __j < __block_size; ++__j) {
0092       __init = __binary_op(std::move(__init), std::move(__lane[__j]));
0093     }
0094     // destroyer
0095     _PSTL_PRAGMA_SIMD
0096     for (_Size __j = 0; __j < __block_size; ++__j) {
0097       __lane[__j].~_Tp();
0098     }
0099   } else {
0100     for (_Size __i = 0; __i < __n; ++__i) {
0101       __init = __binary_op(std::move(__init), __f(__i));
0102     }
0103   }
0104   return __init;
0105 }
0106
0107 template <class _Backend, class _RawExecutionPolicy>
0108 struct __cpu_parallel_transform_reduce_binary {
0109   template <class _Policy,
0110             class _ForwardIterator1,
0111             class _ForwardIterator2,
0112             class _Tp,
0113             class _BinaryOperation1,
0114             class _BinaryOperation2>
0115   _LIBCPP_HIDE_FROM_ABI optional<_Tp> operator()(
0116       _Policy&& __policy,
0117       _ForwardIterator1 __first1,
0118       _ForwardIterator1 __last1,
0119       _ForwardIterator2 __first2,
0120       _Tp __init,
0121       _BinaryOperation1 __reduce,
0122       _BinaryOperation2 __transform) const noexcept {
0123     if constexpr (__is_parallel_execution_policy_v<_RawExecutionPolicy> &&
0124                   __has_random_access_iterator_category_or_concept<_ForwardIterator1>::value &&
0125                   __has_random_access_iterator_category_or_concept<_ForwardIterator2>::value) {
0126       return __cpu_traits<_Backend>::__transform_reduce(
0127           __first1,
0128           std::move(__last1),
0129           [__first1, __first2, __transform](_ForwardIterator1 __iter) {
0130             return __transform(*__iter, *(__first2 + (__iter - __first1)));
0131           },
0132           std::move(__init),
0133           std::move(__reduce),
0134           [&__policy, __first1, __first2, __reduce, __transform](
0135               _ForwardIterator1 __brick_first, _ForwardIterator1 __brick_last, _Tp __brick_init) {
0136             using _TransformReduceBinaryUnseq =
0137                 __pstl::__transform_reduce_binary<_Backend, __remove_parallel_policy_t<_RawExecutionPolicy>>;
0138             return *_TransformReduceBinaryUnseq()(
0139                 std::__remove_parallel_policy(__policy),
0140                 __brick_first,
0141                 std::move(__brick_last),
0142                 __first2 + (__brick_first - __first1),
0143                 std::move(__brick_init),
0144                 std::move(__reduce),
0145                 std::move(__transform));
0146           });
0147     } else if constexpr (__is_unsequenced_execution_policy_v<_RawExecutionPolicy> &&
0148                          __has_random_access_iterator_category_or_concept<_ForwardIterator1>::value &&
0149                          __has_random_access_iterator_category_or_concept<_ForwardIterator2>::value) {
0150       return __pstl::__simd_transform_reduce<_Backend>(
0151           __last1 - __first1, std::move(__init), std::move(__reduce), [&](__iter_diff_t<_ForwardIterator1> __i) {
0152             return __transform(__first1[__i], __first2[__i]);
0153           });
0154     } else {
0155       return std::transform_reduce(
0156           std::move(__first1),
0157           std::move(__last1),
0158           std::move(__first2),
0159           std::move(__init),
0160           std::move(__reduce),
0161           std::move(__transform));
0162     }
0163   }
0164 };
0165
0166 template <class _Backend, class _RawExecutionPolicy>
0167 struct __cpu_parallel_transform_reduce {
0168   template <class _Policy, class _ForwardIterator, class _Tp, class _BinaryOperation, class _UnaryOperation>
0169   _LIBCPP_HIDE_FROM_ABI optional<_Tp>
0170   operator()(_Policy&& __policy,
0171              _ForwardIterator __first,
0172              _ForwardIterator __last,
0173              _Tp __init,
0174              _BinaryOperation __reduce,
0175              _UnaryOperation __transform) const noexcept {
0176     if constexpr (__is_parallel_execution_policy_v<_RawExecutionPolicy> &&
0177                   __has_random_access_iterator_category_or_concept<_ForwardIterator>::value) {
0178       return __cpu_traits<_Backend>::__transform_reduce(
0179           std::move(__first),
0180           std::move(__last),
0181           [__transform](_ForwardIterator __iter) { return __transform(*__iter); },
0182           std::move(__init),
0183           __reduce,
0184           [&__policy, __transform, __reduce](auto __brick_first, auto __brick_last, _Tp __brick_init) {
0185             using _TransformReduceUnseq =
0186                 __pstl::__transform_reduce<_Backend, __remove_parallel_policy_t<_RawExecutionPolicy>>;
0187             auto __res = _TransformReduceUnseq()(
0188                 std::__remove_parallel_policy(__policy),
0189                 std::move(__brick_first),
0190                 std::move(__brick_last),
0191                 std::move(__brick_init),
0192                 std::move(__reduce),
0193                 std::move(__transform));
0194             _LIBCPP_ASSERT_INTERNAL(__res, "unseq/seq should never try to allocate!");
0195             return *std::move(__res);
0196           });
0197     } else if constexpr (__is_unsequenced_execution_policy_v<_RawExecutionPolicy> &&
0198                          __has_random_access_iterator_category_or_concept<_ForwardIterator>::value) {
0199       return __pstl::__simd_transform_reduce<_Backend>(
0200           __last - __first,
0201           std::move(__init),
0202           std::move(__reduce),
0203           [=, &__transform](__iter_diff_t<_ForwardIterator> __i) { return __transform(__first[__i]); });
0204     } else {
0205       return std::transform_reduce(
0206           std::move(__first), std::move(__last), std::move(__init), std::move(__reduce), std::move(__transform));
0207     }
0208   }
0209 };
0210
0211 } // namespace __pstl
0212 _LIBCPP_END_NAMESPACE_STD
0213
0214 #endif // _LIBCPP_STD_VER >= 17
0215
0216 _LIBCPP_POP_MACROS
0217
0218 #endif // _LIBCPP___PSTL_CPU_ALGOS_TRANSFORM_REDUCE_H