Vectorize search_n for small values of n

benchmarks/src/search_n.cpp

@@ -5,52 +5,87 @@
 #include <benchmark/benchmark.h>
 #include <cstddef>
 #include <cstdint>
+#include <random>
 #include <vector>
 
 #include "skewed_allocator.hpp"
 
 using namespace std;
 
-// NB: This particular algorithm has std and ranges implementations with different perf characteristics!
+// NB: This particular algorithm has std and ranges non-vectorized implementations with different perf characteristics!
 
 enum class AlgType { Std, Rng };
 
-template <class T, AlgType Alg>
+enum class PatternType {
+    TwoZones,
+    DenseSmallSequences,
+};
+
+template <class T, AlgType Alg, PatternType Pattern>
 void bm(benchmark::State& state) {
     const auto size = static_cast<size_t>(state.range(0));
-
-    constexpr size_t N = 1;
+    const auto n    = static_cast<size_t>(state.range(1));
 
     constexpr T no_match{'-'};
     constexpr T match{'*'};
 
     vector<T, not_highly_aligned_allocator<T>> v(size, no_match);
 
-    fill(v.begin() + v.size() / 2, v.end(), match);
+    if constexpr (Pattern == PatternType::TwoZones) {
+        fill(v.begin() + v.size() / 2, v.end(), match);
+    } else if constexpr (Pattern == PatternType::DenseSmallSequences) {
+        if (size != 0 && n != 0) {
+            mt19937 gen{7687239};
+
+            uniform_int_distribution<size_t> len_dis(0, n - 1);
+
+            size_t cur_len = len_dis(gen);
+
+            for (size_t i = 0; i != size; ++i) {
+                if (cur_len != 0) {
+                    v[i] = match;
+                    --cur_len;
+                } else {
+                    cur_len = len_dis(gen);
+                }
+            }
+        }
+    }
 
     for (auto _ : state) {
         if constexpr (Alg == AlgType::Std) {
-            benchmark::DoNotOptimize(search_n(v.begin(), v.end(), N, match));
+            benchmark::DoNotOptimize(search_n(v.begin(), v.end(), n, match));
         } else if constexpr (Alg == AlgType::Rng) {
-            benchmark::DoNotOptimize(ranges::search_n(v, N, match));
+            benchmark::DoNotOptimize(ranges::search_n(v, n, match));
         }
     }
 }
 
 void common_args(auto bm) {
-    bm->Arg(3000);
+    for (const auto& n : {40, 18, 16, 14, 10, 8, 5, 4, 3, 2, 1}) {
+        bm->ArgPair(3000, n);
+    }
 }
 
-BENCHMARK(bm<uint8_t, AlgType::Std>)->Apply(common_args);
-BENCHMARK(bm<uint8_t, AlgType::Rng>)->Apply(common_args);
+BENCHMARK(bm<uint8_t, AlgType::Std, PatternType::TwoZones>)->Apply(common_args);
+BENCHMARK(bm<uint8_t, AlgType::Rng, PatternType::TwoZones>)->Apply(common_args);
+BENCHMARK(bm<uint8_t, AlgType::Std, PatternType::DenseSmallSequences>)->Apply(common_args);
+BENCHMARK(bm<uint8_t, AlgType::Rng, PatternType::DenseSmallSequences>)->Apply(common_args);
+
+BENCHMARK(bm<uint16_t, AlgType::Std, PatternType::TwoZones>)->Apply(common_args);
+BENCHMARK(bm<uint16_t, AlgType::Rng, PatternType::TwoZones>)->Apply(common_args);
+BENCHMARK(bm<uint16_t, AlgType::Std, PatternType::DenseSmallSequences>)->Apply(common_args);
+BENCHMARK(bm<uint16_t, AlgType::Rng, PatternType::DenseSmallSequences>)->Apply(common_args);
 
-BENCHMARK(bm<uint16_t, AlgType::Std>)->Apply(common_args);
-BENCHMARK(bm<uint16_t, AlgType::Rng>)->Apply(common_args);
+BENCHMARK(bm<uint32_t, AlgType::Std, PatternType::TwoZones>)->Apply(common_args);
+BENCHMARK(bm<uint32_t, AlgType::Rng, PatternType::TwoZones>)->Apply(common_args);
+BENCHMARK(bm<uint32_t, AlgType::Std, PatternType::DenseSmallSequences>)->Apply(common_args);
+BENCHMARK(bm<uint32_t, AlgType::Rng, PatternType::DenseSmallSequences>)->Apply(common_args);
 
-BENCHMARK(bm<uint32_t, AlgType::Std>)->Apply(common_args);
-BENCHMARK(bm<uint32_t, AlgType::Rng>)->Apply(common_args);
+BENCHMARK(bm<uint64_t, AlgType::Std, PatternType::TwoZones>)->Apply(common_args);
+BENCHMARK(bm<uint64_t, AlgType::Rng, PatternType::TwoZones>)->Apply(common_args);
+BENCHMARK(bm<uint64_t, AlgType::Std, PatternType::DenseSmallSequences>)->Apply(common_args);
+BENCHMARK(bm<uint64_t, AlgType::Rng, PatternType::DenseSmallSequences>)->Apply(common_args);
 
-BENCHMARK(bm<uint64_t, AlgType::Std>)->Apply(common_args);
-BENCHMARK(bm<uint64_t, AlgType::Rng>)->Apply(common_args);
 
 BENCHMARK_MAIN();

stl/inc/algorithm

@@ -80,6 +80,11 @@ __declspec(noalias) void __stdcall __std_replace_4(
 __declspec(noalias) void __stdcall __std_replace_8(
     void* _First, void* _Last, uint64_t _Old_val, uint64_t _New_val) noexcept;
 
+const void* __stdcall __std_search_n_1(const void* _First, const void* _Last, size_t _Count, uint8_t _Value) noexcept;
+const void* __stdcall __std_search_n_2(const void* _First, const void* _Last, size_t _Count, uint16_t _Value) noexcept;
+const void* __stdcall __std_search_n_4(const void* _First, const void* _Last, size_t _Count, uint32_t _Value) noexcept;
+const void* __stdcall __std_search_n_8(const void* _First, const void* _Last, size_t _Count, uint64_t _Value) noexcept;
+
 void* __stdcall __std_unique_1(void* _First, void* _Last) noexcept;
 void* __stdcall __std_unique_2(void* _First, void* _Last) noexcept;
 void* __stdcall __std_unique_4(void* _First, void* _Last) noexcept;
@@ -210,6 +215,33 @@ __declspec(noalias) void _Replace_vectorized(
     }
 }
 
+template <class _Ty, class _TVal>
+_Ty* _Search_n_vectorized(_Ty* const _First, _Ty* const _Last, const size_t _Count, const _TVal _Val) noexcept {
+    if constexpr (is_pointer_v<_Ty>) {
+#ifdef _WIN64
+        return const_cast<_Ty*>(
+            static_cast<const _Ty*>(::__std_search_n_8(_First, _Last, _Count, reinterpret_cast<uint64_t>(_Val))));
+#else // ^^^ defined(_WIN64) / !defined(_WIN64) vvv
+        return const_cast<_Ty*>(
+            static_cast<const _Ty*>(::__std_search_n_4(_First, _Last, _Count, reinterpret_cast<uint32_t>(_Val))));
+#endif // ^^^ !defined(_WIN64) ^^^
+    } else if constexpr (sizeof(_Ty) == 1) {
+        return const_cast<_Ty*>(
+            static_cast<const _Ty*>(::__std_search_n_1(_First, _Last, _Count, static_cast<uint8_t>(_Val))));
+    } else if constexpr (sizeof(_Ty) == 2) {
+        return const_cast<_Ty*>(
+            static_cast<const _Ty*>(::__std_search_n_2(_First, _Last, _Count, static_cast<uint16_t>(_Val))));
+    } else if constexpr (sizeof(_Ty) == 4) {
+        return const_cast<_Ty*>(
+            static_cast<const _Ty*>(::__std_search_n_4(_First, _Last, _Count, static_cast<uint32_t>(_Val))));
+    } else if constexpr (sizeof(_Ty) == 8) {
+        return const_cast<_Ty*>(
+            static_cast<const _Ty*>(::__std_search_n_8(_First, _Last, _Count, static_cast<uint64_t>(_Val))));
+    } else {
+        _STL_INTERNAL_STATIC_ASSERT(false); // unexpected size
+    }
+}
+
 template <class _Ty>
 _Ty* _Unique_vectorized(_Ty* const _First, _Ty* const _Last) noexcept {
     if constexpr (sizeof(_Ty) == 1) {
@@ -240,6 +272,13 @@ constexpr bool _Vector_alg_in_ranges_replace_is_safe =
     _Vector_alg_in_replace_is_safe<_Iter, _Ty1> // can search and replace
     && _Vector_alg_in_find_is_safe_elem<_Ty2, _Iter_value_t<_Iter>>; // replacement fits
 
+template <class _Iter, class _Ty, class _Pr>
+constexpr bool _Vector_alg_in_search_n_is_safe = _Vector_alg_in_find_is_safe<_Iter, _Ty>
+                                              && _Is_any_of_v<_Pr,
+#if _HAS_CXX20
+                                                  ranges::equal_to,
+#endif
+                                                  equal_to<>>;
 // Can we activate the vector algorithms for unique?
 template <class _Iter, class _Pr>
 constexpr bool _Vector_alg_in_unique_is_safe = _Equal_memcmp_is_safe<_Iter, _Iter, _Pr>;
@@ -2281,6 +2320,29 @@ _NODISCARD _CONSTEXPR20 _FwdIt search_n(
     auto _UFirst      = _STD _Get_unwrapped(_First);
     const auto _ULast = _STD _Get_unwrapped(_Last);
     if constexpr (_Is_ranges_random_iter_v<_FwdIt>) {
+#if _USE_STD_VECTOR_ALGORITHMS
+        if constexpr (_Vector_alg_in_search_n_is_safe<decltype(_UFirst), _Ty, _Pr>) {
+            if (!_STD _Is_constant_evaluated()) {
+                if (!_STD _Could_compare_equal_to_value_type<decltype(_UFirst)>(_Val)) {
+                    return _Last;
+                }
+
+                const auto _First_ptr = _STD _To_address(_UFirst);
+                const auto _Result =
+                    _STD _Search_n_vectorized(_First_ptr, _STD _To_address(_ULast), static_cast<size_t>(_Count), _Val);
+
+                if constexpr (is_pointer_v<decltype(_UFirst)>) {
+                    _UFirst = _Result;
+                } else {
+                    _UFirst += _Result - _First_ptr;
+                }
+
+                _STD _Seek_wrapped(_Last, _UFirst);
+                return _Last;
+            }
+        }
+#endif // _USE_STD_VECTOR_ALGORITHMS
+
         const auto _Count_diff = static_cast<_Iter_diff_t<_FwdIt>>(_Count);
         auto _UOld_first       = _UFirst;
         for (_Iter_diff_t<_FwdIt> _Inc = 0; _Count_diff <= _ULast - _UOld_first;) { // enough room, look for a match
@@ -2453,6 +2515,38 @@ namespace ranges {
                     return {_First, _First};
                 }
 
+#if _USE_STD_VECTOR_ALGORITHMS
+                if constexpr (_Vector_alg_in_search_n_is_safe<_It, _Ty, _Pr> && is_same_v<_Pj, identity>) {
+                    if (!_STD is_constant_evaluated()) {
+                        if (!_STD _Could_compare_equal_to_value_type<_It>(_Val)) {
+                            _First += _Dist;
+                            return {_First, _First};
+                        }
+
+                        const auto _First_ptr = _STD _To_address(_First);
+                        const auto _Last_ptr  = _First_ptr + _Dist;
+                        const auto _Result =
+                            _STD _Search_n_vectorized(_First_ptr, _Last_ptr, static_cast<size_t>(_Count), _Val);
+
+                        if constexpr (is_pointer_v<_It>) {
+                            if (_Result != _Last_ptr) {
+                                return {_Result, _Result + _Count};
+                            } else {
+                                return {_Result, _Result};
+                            }
+                        } else {
+                            if (_Result != _Last_ptr) {
+                                _First += _Result - _First_ptr;
+                                return {_First, _First + _Count};
+                            } else {
+                                _First += _Dist;
+                                return {_First, _First};
+                            }
+                        }
+                    }
+                }
+#endif // _USE_STD_VECTOR_ALGORITHMS
+
                 auto _Last = _RANGES next(_First, _Count);
                 auto _Mid1 = _First;
                 auto _Mid2 = _Last;

stl/src/vector_algorithms.cpp

@@ -2563,7 +2563,7 @@ namespace {
     // In optimized builds it avoids an extra call, as these functions are too large to inline.
 
     template <class _Traits, _Find_one_predicate _Pred, class _Ty>
-    const void* __stdcall __std_find_trivial_impl(const void* _First, const void* _Last, _Ty _Val) noexcept {
+    const void* __stdcall __std_find_trivial_impl(const void* _First, const void* const _Last, _Ty _Val) noexcept {
 #ifndef _M_ARM64EC
         const size_t _Size_bytes = _Byte_length(_First, _Last);
 
@@ -3066,6 +3066,124 @@ namespace {
         return _Result;
     }
 
+    template <class _Traits, class _Ty>
+    const void* __stdcall __std_search_n_impl(
+        const void* _First, const void* const _Last, const size_t _Count, const _Ty _Val) noexcept {
+        if (_Count == 0) {
+            return _First;
+        } else if (_Count == 1) {
+            return __std_find_trivial_impl<_Traits, _Find_one_predicate::_Equal>(_First, _Last, _Val);
+        }
+
+        auto _Mid1 = static_cast<const _Ty*>(_First);
+#ifndef _M_ARM64EC
+        const size_t _Length = _Byte_length(_First, _Last);
+        if (_Count <= 16 / sizeof(_Ty) && _Length >= 32 && _Use_avx2()) {
+            _Zeroupper_on_exit _Guard; // TRANSITION, DevCom-10331414
+
+            const int _Bytes_count = static_cast<int>(_Count * sizeof(_Ty));
+            const int _Sh1         = sizeof(_Ty) == 1 ? (_Bytes_count < 4 ? _Bytes_count - 2 : 2) : 0;
+            const int _Sh2 = sizeof(_Ty) < 4 ? (_Bytes_count < 4 ? 0 : (_Bytes_count < 8 ? _Bytes_count - 4 : 4)) : 0;
+            const int _Sh3 = sizeof(_Ty) < 8 ? (_Bytes_count < 8 ? 0 : _Bytes_count - 8) : 0;
+
+            const __m256i _Comparand = _Traits::_Set_avx(_Val);
+
+            const void* _Stop_at = _First;
+            _Advance_bytes(_Stop_at, _Length & ~size_t{0x1F});
+
+            uint32_t _Carry = 0;
+            do {
+                const __m256i _Data = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(_First));
+
+                const auto _Mask = static_cast<uint32_t>(_mm256_movemask_epi8(_Traits::_Cmp_avx(_Comparand, _Data)));
+
+                uint64_t _MskX = uint64_t{_Carry} | (uint64_t{_Mask} << 32);
+
+                if constexpr (sizeof(_Ty) == 1) {
+                    _MskX = (_MskX >> 1) & _MskX;
+                    _MskX = __ull_rshift(_MskX, _Sh1) & _MskX;
+                }
+
+                if constexpr (sizeof(_Ty) == 2) {
+                    _MskX = (_MskX >> 2) & _MskX;
+                }
+
+                if constexpr (sizeof(_Ty) < 4) {
+                    _MskX = __ull_rshift(_MskX, _Sh2) & _MskX;
+                }
+
+                if constexpr (sizeof(_Ty) == 4) {
+                    _MskX = (_MskX >> 4) & _MskX;
+                }
+
+                if constexpr (sizeof(_Ty) < 8) {
+                    _MskX = __ull_rshift(_MskX, _Sh3) & _MskX;
+                }
+
+                if constexpr (sizeof(_Ty) == 8) {
+                    _MskX = (_MskX >> 8) & _MskX;
+                }
+
+                if (_MskX != 0) {
+#ifdef _M_IX86
+                    const uint32_t _MskLow = static_cast<uint32_t>(_MskX);
+
+                    const int _Shift = _MskLow != 0 ? static_cast<int>(_tzcnt_u32(_MskLow)) - 32
+                                                    : static_cast<int>(_tzcnt_u32(static_cast<uint32_t>(_MskX >> 32)));
+
+#elifdef _M_X64
+                    const long long _Shift = static_cast<long long>(_tzcnt_u64(_MskX)) - 32;
+#else
+#error Unsupported architecture
+#endif
+                    _Advance_bytes(_First, _Shift);
+                    return _First;
+                }
+
+                _Carry = _Mask;
+
+                _Advance_bytes(_First, 32);
+            } while (_First != _Stop_at);
+
+            _Mid1 = static_cast<const _Ty*>(_First);
+            _Rewind_bytes(_First, _lzcnt_u32(~_Carry));
+        }
+#endif // !_M_ARM64EC
+        auto _Match_start    = static_cast<const _Ty*>(_First);
+        const auto _Last_ptr = static_cast<const _Ty*>(_Last);
+
+        if (static_cast<size_t>(_Last_ptr - _Match_start) < _Count) {
+            return _Last_ptr;
+        }
+
+        auto _Match_end = _Match_start + _Count;
+        auto _Mid2      = _Match_end;
+        for (;;) {
+            // Invariants: _Match_end - _Match_start == _Count, [_Match_start, _Mid1) and [_Mid2, _Match_end) match
+            // _Val:
+            //
+            // _Match_start  _Mid1    _Mid2    _Match_end
+            // |=============|????????|========|??????????...
+
+            --_Mid2;
+            if (*_Mid2 == _Val) { // match;
+                if (_Mid1 == _Mid2) { // [_Mid1, _Mid2) is empty, so [_Match_start, _Match_end) all match
+                    return _Match_start;
+                }
+            } else { // mismatch; skip past it
+                _Match_start = _Mid2 + 1;
+
+                if (static_cast<size_t>(_Last_ptr - _Match_start) < _Count) { // not enough space left
+                    return _Last_ptr;
+                }
+
+                _Mid1      = _Match_end;
+                _Match_end = _Match_start + _Count;
+                _Mid2      = _Match_end;
+            }
+        }
+    }
+
     enum class _Find_meow_of_predicate { _Any_of, _None_of };
 
 #ifndef _M_ARM64EC
@@ -4914,6 +5032,26 @@ __declspec(noalias) size_t __stdcall __std_count_trivial_8(
     return __std_count_trivial_impl<_Count_traits_8>(_First, _Last, _Val);
 }
 
+const void* __stdcall __std_search_n_1(
+    const void* const _First, const void* const _Last, const size_t _Count, const uint8_t _Value) noexcept {
+    return __std_search_n_impl<_Find_traits_1>(_First, _Last, _Count, _Value);
+}
+
+const void* __stdcall __std_search_n_2(
+    const void* const _First, const void* const _Last, const size_t _Count, const uint16_t _Value) noexcept {
+    return __std_search_n_impl<_Find_traits_2>(_First, _Last, _Count, _Value);
+}
+
+const void* __stdcall __std_search_n_4(
+    const void* const _First, const void* const _Last, const size_t _Count, const uint32_t _Value) noexcept {
+    return __std_search_n_impl<_Find_traits_4>(_First, _Last, _Count, _Value);
+}
+
+const void* __stdcall __std_search_n_8(
+    const void* const _First, const void* const _Last, const size_t _Count, const uint64_t _Value) noexcept {
+    return __std_search_n_impl<_Find_traits_8>(_First, _Last, _Count, _Value);
+}
+
 const void* __stdcall __std_find_first_of_trivial_1(
     const void* const _First1, const void* const _Last1, const void* const _First2, const void* const _Last2) noexcept {
     return __std_find_first_of::_Dispatch_ptr<uint8_t>(_First1, _Last1, _First2, _Last2);