From 4e5136ceab6abcc04f76bd9cc97bc019798da7c1 Mon Sep 17 00:00:00 2001
From: Alex Guteniev <gutenev@gmail.com>
Date: Sun, 12 May 2024 12:24:41 +0300
Subject: [PATCH 1/2] auto vectorize `std::count`

---
 benchmarks/src/find_and_count.cpp |  4 +-
 stl/inc/xutility                  | 68 ++++++++++++++++++++++++++++---
 2 files changed, 65 insertions(+), 7 deletions(-)

diff --git a/benchmarks/src/find_and_count.cpp b/benchmarks/src/find_and_count.cpp
index 9c608bfe356..65527403ba9 100644
--- a/benchmarks/src/find_and_count.cpp
+++ b/benchmarks/src/find_and_count.cpp
@@ -33,12 +33,14 @@ void bm(benchmark::State& state) {
     }
 
     for (auto _ : state) {
+        benchmark::DoNotOptimize(a);
+
         if constexpr (Operation == Op::FindSized) {
             benchmark::DoNotOptimize(ranges::find(a.begin(), a.end(), T{'1'}));
         } else if constexpr (Operation == Op::FindUnsized) {
             benchmark::DoNotOptimize(ranges::find(a.begin(), unreachable_sentinel, T{'1'}));
         } else if constexpr (Operation == Op::Count) {
-            benchmark::DoNotOptimize(ranges::count(a.begin(), a.end(), T{'1'}));
+            benchmark::DoNotOptimize(count(a.begin(), a.end(), T{'1'}));
         }
     }
 }
diff --git a/stl/inc/xutility b/stl/inc/xutility
index 0fc9b352a0c..adc1abbe464 100644
--- a/stl/inc/xutility
+++ b/stl/inc/xutility
@@ -5975,6 +5975,22 @@ _NODISCARD constexpr bool _Could_compare_equal_to_value_type(const _Ty& _Val) {
     }
 }
 
+template <class _Ty>
+using _Auto_vectorize_counter = //
+    conditional_t<sizeof(_Ty) == 8, unsigned long long, //
+        conditional_t<sizeof(_Ty) == 4, unsigned int, //
+            conditional_t<sizeof(_Ty) == 2, unsigned short, //
+                conditional_t<sizeof(_Ty) == 1, unsigned char, void>>>>;
+
+
+// Can we attempt auto vectorization for count?
+template <class _Iter>
+constexpr bool _Can_try_auto_vectorize_count =
+    _Iterator_is_contiguous<_Iter> // The iterator must be contiguous
+    && !_Iterator_is_volatile<_Iter> // volatile iterators do ono auto vectorize
+    && is_fundamental_v<_Iter_value_t<_Iter>> // MSVC only auto vectorizes counting of fundametal types
+    && !is_same_v<_Auto_vectorize_counter<_Iter_value_t<_Iter>>, void>; // Vector element size matches
+
 template <class _InIt, class _Ty>
 _NODISCARD _CONSTEXPR20 _InIt _Find_unchecked(_InIt _First, const _InIt _Last, const _Ty& _Val) {
     // find first matching _Val; choose optimization
@@ -6170,16 +6186,56 @@ _NODISCARD _CONSTEXPR20 _Iter_diff_t<_InIt> count(const _InIt _First, const _InI
             }
         }
 #endif // _USE_STD_VECTOR_ALGORITHMS
+        using _Result_type = _Iter_diff_t<_InIt>;
 
-        _Iter_diff_t<_InIt> _Count = 0;
+        if constexpr (_Can_try_auto_vectorize_count<decltype(_UFirst)>) {
+            using _Counter_type = _Auto_vectorize_counter<_Iter_value_t<decltype(_UFirst)>>;
 
-        for (; _UFirst != _ULast; ++_UFirst) {
-            if (*_UFirst == _Val) {
-                ++_Count;
+            if constexpr (sizeof(_Counter_type) >= sizeof(_Result_type)) {
+                _Counter_type _Count = 0;
+
+                for (; _UFirst != _ULast; ++_UFirst) {
+                    if (*_UFirst == _Val) {
+                        ++_Count;
+                    }
+                }
+
+                return static_cast<_Result_type>(_Count);
+            } else {
+                _Result_type _Outer_count                 = 0;
+                constexpr _Counter_type _Max_portion_size = _Max_limit<_Counter_type>();
+
+                while (_UFirst != _ULast) {
+                    _Counter_type _Inner_count = 0;
+                    _Result_type _Portion_size = static_cast<_Result_type>(_ULast - _UFirst);
+                    if (_Portion_size > _Max_portion_size) {
+                        _Portion_size = _Max_portion_size;
+                    }
+
+                    const auto _UStop = _UFirst + _Portion_size;
+
+                    for (; _UFirst != _UStop; ++_UFirst) {
+                        if (*_UFirst == _Val) {
+                            ++_Inner_count;
+                        }
+                    }
+
+                    _Outer_count += _Inner_count;
+                }
+
+                return _Outer_count;
+            }
+        } else {
+            _Result_type _Count = 0;
+
+            for (; _UFirst != _ULast; ++_UFirst) {
+                if (*_UFirst == _Val) {
+                    ++_Count;
+                }
             }
-        }
 
-        return _Count;
+            return _Count;
+        }
     }
 }
 

From ae43611e9deea81cf16ac0b2983aa97f732392d8 Mon Sep 17 00:00:00 2001
From: Alex Guteniev <gutenev@gmail.com>
Date: Sun, 12 May 2024 14:01:58 +0300
Subject: [PATCH 2/2] avoid unvectorized tail for portions

---
 stl/inc/xutility | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/stl/inc/xutility b/stl/inc/xutility
index adc1abbe464..800ac9bd538 100644
--- a/stl/inc/xutility
+++ b/stl/inc/xutility
@@ -6203,7 +6203,7 @@ _NODISCARD _CONSTEXPR20 _Iter_diff_t<_InIt> count(const _InIt _First, const _InI
                 return static_cast<_Result_type>(_Count);
             } else {
                 _Result_type _Outer_count                 = 0;
-                constexpr _Counter_type _Max_portion_size = _Max_limit<_Counter_type>();
+                constexpr _Counter_type _Max_portion_size = _Counter_type{1} << (sizeof(_Counter_type) * 8 - 1);
 
                 while (_UFirst != _ULast) {
                     _Counter_type _Inner_count = 0;