[libcxx-commits] [libcxx] [libc++] Speed-up the reduction by up to 4X by breaking the dependency chain (PR #234887)

[via libcxx-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-libcxx

Author: Michael G. Kazakov (mikekazakov)

<details>
<summary>Changes</summary>

This PR changes `std::reduce`/`std::transform_reduce()` to process 4 elements per iteration and perform reduction in a tree-like manner:
```
__init--------------\
                    (X) -> __init
__first[0]\         /
          (X)\     /
__first[1]/   \   /
              (X)/
__first[2]\   /
          (X)/
__first[3]/
```
This breaks up the dependency chain to allow the CPU to perform more speculative execution and to utilise more execution units.
On Apple M5 this leads to ~4X speedup in the tested benchmarks.
Switching to 8-way reduction brings up to ~6X speedup in some benchmarks, but can also degrade benchmarks with trivial operators where previously the compiler could vectorise the computation, so the 4-way reduction is chosen as a safe option.
Of course, in the case of very compute-costly reduction operators, the effect of this manual unrolling would be lower.
This optimisation applies only when random-access iterators are passed into `std::reduce`/`std::transform_reduce()`.

---
Full diff: https://github.com/llvm/llvm-project/pull/234887.diff


2 Files Affected:

- (modified) libcxx/include/__numeric/reduce.h (+6-2) 
- (modified) libcxx/include/__numeric/transform_reduce.h (+15-2) 


``````````diff
diff --git a/libcxx/include/__numeric/reduce.h b/libcxx/include/__numeric/reduce.h
index 853ca6fb61860..fe245e9097801 100644
--- a/libcxx/include/__numeric/reduce.h
+++ b/libcxx/include/__numeric/reduce.h
@@ -27,8 +27,12 @@ _LIBCPP_PUSH_MACROS
 _LIBCPP_BEGIN_NAMESPACE_STD
 
 template <class _InputIterator, class _Tp, class _BinaryOp>
-[[nodiscard]] _LIBCPP_HIDE_FROM_ABI
-_LIBCPP_CONSTEXPR_SINCE_CXX20 _Tp reduce(_InputIterator __first, _InputIterator __last, _Tp __init, _BinaryOp __b) {
+[[nodiscard]] _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 _Tp
+reduce(_InputIterator __first, _InputIterator __last, _Tp __init, _BinaryOp __b) {
+  if constexpr (__has_random_access_iterator_category<_InputIterator>::value) {
+    for (; __last - __first >= 4; __first += 4)
+      __init = __b(std::move(__init), __b(__b(__first[0], __first[1]), __b(__first[2], __first[3])));
+  }
   for (; __first != __last; ++__first)
     __init = __b(std::move(__init), *__first);
   return __init;
diff --git a/libcxx/include/__numeric/transform_reduce.h b/libcxx/include/__numeric/transform_reduce.h
index 53ed6aaf32d1e..66428b328ae5f 100644
--- a/libcxx/include/__numeric/transform_reduce.h
+++ b/libcxx/include/__numeric/transform_reduce.h
@@ -12,6 +12,7 @@
 
 #include <__config>
 #include <__functional/operations.h>
+#include <__iterator/iterator_traits.h>
 #include <__utility/move.h>
 
 #if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)
@@ -26,8 +27,13 @@ _LIBCPP_PUSH_MACROS
 _LIBCPP_BEGIN_NAMESPACE_STD
 
 template <class _InputIterator, class _Tp, class _BinaryOp, class _UnaryOp>
-[[nodiscard]] _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20
-_Tp transform_reduce(_InputIterator __first, _InputIterator __last, _Tp __init, _BinaryOp __b, _UnaryOp __u) {
+[[nodiscard]] _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20 _Tp
+transform_reduce(_InputIterator __first, _InputIterator __last, _Tp __init, _BinaryOp __b, _UnaryOp __u) {
+  if constexpr (__has_random_access_iterator_category<_InputIterator>::value) {
+    for (; __last - __first >= 4; __first += 4)
+      __init =
+          __b(std::move(__init), __b(__b(__u(__first[0]), __u(__first[1])), __b(__u(__first[2]), __u(__first[3]))));
+  }
   for (; __first != __last; ++__first)
     __init = __b(std::move(__init), __u(*__first));
   return __init;
@@ -41,6 +47,13 @@ template <class _InputIterator1, class _InputIterator2, class _Tp, class _Binary
     _Tp __init,
     _BinaryOp1 __b1,
     _BinaryOp2 __b2) {
+  if constexpr (__has_random_access_iterator_category<_InputIterator1>::value &&
+                __has_random_access_iterator_category<_InputIterator2>::value) {
+    for (; __last1 - __first1 >= 4; __first1 += 4, (void)(__first2 += 4))
+      __init = __b1(std::move(__init),
+                    __b1(__b1(__b2(__first1[0], __first2[0]), __b2(__first1[1], __first2[1])),
+                         __b1(__b2(__first1[2], __first2[2]), __b2(__first1[3], __first2[3]))));
+  }
   for (; __first1 != __last1; ++__first1, (void)++__first2)
     __init = __b1(std::move(__init), __b2(*__first1, *__first2));
   return __init;

``````````

</details>


https://github.com/llvm/llvm-project/pull/234887