Improve basic_string::find_first_of and basic_string::find_last_of vectorization for large needles or very large haystacks (#5029)

AlexGuteniev · StephanTLavavej · web-flow · commit eaf73552f64b · 2024-12-12T22:06:41.000-08:00
Co-authored-by: Stephan T. Lavavej &lt;stl@nuwen.net&gt;
diff --git a/benchmarks/src/find_first_of.cpp b/benchmarks/src/find_first_of.cpp
@@ -13,6 +13,8 @@
 #include <type_traits>
 #include <vector>
 
+#include "skewed_allocator.hpp"
+
 using namespace std;
 
 enum class AlgType { std_func, str_member_first, str_member_last };
@@ -24,7 +26,8 @@ void bm(benchmark::State& state) {
     const size_t HSize = Pos * 2;
     const size_t Which = 0;
 
-    using container = conditional_t<Alg == AlgType::std_func, vector<T>, basic_string<T>>;
+    using container = conditional_t<Alg == AlgType::std_func, vector<T, not_highly_aligned_allocator<T>>,
+        basic_string<T, char_traits<T>, not_highly_aligned_allocator<T>>>;
 
     constexpr T HaystackFiller{' '};
     static_assert(HaystackFiller < Start, "The following iota() should not produce the haystack filler.");
@@ -59,8 +62,9 @@ void bm(benchmark::State& state) {
 }
 
 void common_args(auto bm) {
-    bm->Args({2, 3})->Args({7, 4})->Args({9, 3})->Args({22, 5})->Args({58, 2})->Args({102, 4});
-    bm->Args({325, 1})->Args({400, 50})->Args({1011, 11})->Args({1502, 23})->Args({3056, 7});
+    bm->Args({2, 3})->Args({6, 81})->Args({7, 4})->Args({9, 3})->Args({22, 5})->Args({58, 2});
+    bm->Args({75, 85})->Args({102, 4})->Args({200, 46})->Args({325, 1})->Args({400, 50});
+    bm->Args({1011, 11})->Args({1280, 46})->Args({1502, 23})->Args({2203, 54})->Args({3056, 7});
 }
 
 BENCHMARK(bm<AlgType::std_func, uint8_t>)->Apply(common_args);
diff --git a/stl/inc/__msvc_string_view.hpp b/stl/inc/__msvc_string_view.hpp
@@ -29,6 +29,15 @@ extern "C" {
 // compiler has to assume that the denoted arrays are "globally address taken", and that any later calls to
 // unanalyzable routines may modify those arrays.
 
+__declspec(noalias) size_t __stdcall __std_find_first_of_trivial_pos_1(
+    const void* _Haystack, size_t _Haystack_length, const void* _Needle, size_t _Needle_length) noexcept;
+__declspec(noalias) size_t __stdcall __std_find_first_of_trivial_pos_2(
+    const void* _Haystack, size_t _Haystack_length, const void* _Needle, size_t _Needle_length) noexcept;
+__declspec(noalias) size_t __stdcall __std_find_first_of_trivial_pos_4(
+    const void* _Haystack, size_t _Haystack_length, const void* _Needle, size_t _Needle_length) noexcept;
+__declspec(noalias) size_t __stdcall __std_find_first_of_trivial_pos_8(
+    const void* _Haystack, size_t _Haystack_length, const void* _Needle, size_t _Needle_length) noexcept;
+
 __declspec(noalias) size_t __stdcall __std_find_last_of_trivial_pos_1(
     const void* _Haystack, size_t _Haystack_length, const void* _Needle, size_t _Needle_length) noexcept;
 __declspec(noalias) size_t __stdcall __std_find_last_of_trivial_pos_2(
@@ -38,6 +47,23 @@ __declspec(noalias) size_t __stdcall __std_find_last_of_trivial_pos_2(
 
 _STD_BEGIN
 
+template <class _Ty1, class _Ty2>
+size_t _Find_first_of_pos_vectorized(const _Ty1* const _Haystack, const size_t _Haystack_length,
+    const _Ty2* const _Needle, const size_t _Needle_length) noexcept {
+    _STL_INTERNAL_STATIC_ASSERT(sizeof(_Ty1) == sizeof(_Ty2));
+    if constexpr (sizeof(_Ty1) == 1) {
+        return ::__std_find_first_of_trivial_pos_1(_Haystack, _Haystack_length, _Needle, _Needle_length);
+    } else if constexpr (sizeof(_Ty1) == 2) {
+        return ::__std_find_first_of_trivial_pos_2(_Haystack, _Haystack_length, _Needle, _Needle_length);
+    } else if constexpr (sizeof(_Ty1) == 4) {
+        return ::__std_find_first_of_trivial_pos_4(_Haystack, _Haystack_length, _Needle, _Needle_length);
+    } else if constexpr (sizeof(_Ty1) == 8) {
+        return ::__std_find_first_of_trivial_pos_8(_Haystack, _Haystack_length, _Needle, _Needle_length);
+    } else {
+        _STL_INTERNAL_STATIC_ASSERT(false); // unexpected size
+    }
+}
+
 template <class _Ty1, class _Ty2>
 size_t _Find_last_of_pos_vectorized(const _Ty1* const _Haystack, const size_t _Haystack_length,
     const _Ty2* const _Needle, const size_t _Needle_length) noexcept {
@@ -834,48 +860,31 @@ constexpr size_t _Traits_find_first_of(_In_reads_(_Hay_size) const _Traits_ptr_t
     const auto _Hay_end   = _Haystack + _Hay_size;
 
     if constexpr (_Is_implementation_handled_char_traits<_Traits>) {
-        if (!_STD _Is_constant_evaluated()) {
-            using _Elem = typename _Traits::char_type;
-
 #if _USE_STD_VECTOR_ALGORITHMS
-            const bool _Try_vectorize = _Hay_size - _Start_at > _Threshold_find_first_of;
-
-            // Additional condition for when the vectorization outperforms the table lookup
-            constexpr size_t _Find_first_of_bitmap_threshold = sizeof(_Elem) == 1 ? 48 : sizeof(_Elem) == 8 ? 8 : 16;
-
-            const bool _Use_bitmap = !_Try_vectorize || _Needle_size > _Find_first_of_bitmap_threshold;
-#else // ^^^ _USE_STD_VECTOR_ALGORITHMS / !_USE_STD_VECTOR_ALGORITHMS vvv
-            const bool _Use_bitmap = true;
-#endif // ^^^ !_USE_STD_VECTOR_ALGORITHMS ^^^
-
-            if (_Use_bitmap) {
-                _String_bitmap<_Elem> _Matches;
-
-                if (_Matches._Mark(_Needle, _Needle + _Needle_size)) {
-                    for (auto _Match_try = _Hay_start; _Match_try < _Hay_end; ++_Match_try) {
-                        if (_Matches._Match(*_Match_try)) {
-                            return static_cast<size_t>(_Match_try - _Haystack); // found a match
-                        }
-                    }
-                    return static_cast<size_t>(-1); // no match
+        if (!_STD _Is_constant_evaluated()) {
+            const size_t _Remaining_size = _Hay_size - _Start_at;
+            if (_Remaining_size + _Needle_size >= _Threshold_find_first_of) {
+                size_t _Pos = _Find_first_of_pos_vectorized(_Hay_start, _Remaining_size, _Needle, _Needle_size);
+                if (_Pos != static_cast<size_t>(-1)) {
+                    _Pos += _Start_at;
                 }
-
-                // couldn't put one of the characters into the bitmap, fall back to vectorized or serial algorithms
+                return _Pos;
             }
+        }
+#endif // _USE_STD_VECTOR_ALGORITHMS
 
-#if _USE_STD_VECTOR_ALGORITHMS
-            if (_Try_vectorize) {
-                const _Traits_ptr_t<_Traits> _Found =
-                    _STD _Find_first_of_vectorized(_Hay_start, _Hay_end, _Needle, _Needle + _Needle_size);
-
-                if (_Found != _Hay_end) {
-                    return static_cast<size_t>(_Found - _Haystack); // found a match
-                } else {
-                    return static_cast<size_t>(-1); // no match
+        _String_bitmap<typename _Traits::char_type> _Matches;
+
+        if (_Matches._Mark(_Needle, _Needle + _Needle_size)) {
+            for (auto _Match_try = _Hay_start; _Match_try < _Hay_end; ++_Match_try) {
+                if (_Matches._Match(*_Match_try)) {
+                    return static_cast<size_t>(_Match_try - _Haystack); // found a match
                 }
             }
-#endif // _USE_STD_VECTOR_ALGORITHMS
+            return static_cast<size_t>(-1); // no match
         }
+
+        // couldn't put one of the characters into the bitmap, fall back to serial algorithm
     }
 
     for (auto _Match_try = _Hay_start; _Match_try < _Hay_end; ++_Match_try) {
@@ -899,47 +908,32 @@ constexpr size_t _Traits_find_last_of(_In_reads_(_Hay_size) const _Traits_ptr_t<
     const auto _Hay_start = (_STD min)(_Start_at, _Hay_size - 1);
 
     if constexpr (_Is_implementation_handled_char_traits<_Traits>) {
-        if (!_STD _Is_constant_evaluated()) {
-            using _Elem = typename _Traits::char_type;
-
-            bool _Use_bitmap = true;
+        using _Elem = typename _Traits::char_type;
 #if _USE_STD_VECTOR_ALGORITHMS
-            bool _Try_vectorize = false;
-
-            if constexpr (sizeof(_Elem) <= 2) {
-                _Try_vectorize = _Hay_start + 1 > _Threshold_find_first_of;
-                // Additional condition for when the vectorization outperforms the table lookup
-                constexpr size_t _Find_last_of_bitmap_threshold = sizeof(_Elem) == 1 ? 48 : 8;
-
-                _Use_bitmap = !_Try_vectorize || _Needle_size > _Find_last_of_bitmap_threshold;
+        if constexpr (sizeof(_Elem) <= 2) {
+            if (!_STD _Is_constant_evaluated()) {
+                const size_t _Remaining_size = _Hay_start + 1;
+                if (_Remaining_size + _Needle_size >= _Threshold_find_first_of) { // same threshold for first/last
+                    return _Find_last_of_pos_vectorized(_Haystack, _Remaining_size, _Needle, _Needle_size);
+                }
             }
+        }
 #endif // _USE_STD_VECTOR_ALGORITHMS
 
-            if (_Use_bitmap) {
-                _String_bitmap<_Elem> _Matches;
-                if (_Matches._Mark(_Needle, _Needle + _Needle_size)) {
-                    for (auto _Match_try = _Haystack + _Hay_start;; --_Match_try) {
-                        if (_Matches._Match(*_Match_try)) {
-                            return static_cast<size_t>(_Match_try - _Haystack); // found a match
-                        }
-
-                        if (_Match_try == _Haystack) {
-                            return static_cast<size_t>(-1); // at beginning, no more chance for match
-                        }
-                    }
+        _String_bitmap<_Elem> _Matches;
+        if (_Matches._Mark(_Needle, _Needle + _Needle_size)) {
+            for (auto _Match_try = _Haystack + _Hay_start;; --_Match_try) {
+                if (_Matches._Match(*_Match_try)) {
+                    return static_cast<size_t>(_Match_try - _Haystack); // found a match
                 }
 
-                // couldn't put one of the characters into the bitmap, fall back to vectorized or serial algorithms
-            }
-
-#if _USE_STD_VECTOR_ALGORITHMS
-            if constexpr (sizeof(_Elem) <= 2) {
-                if (_Try_vectorize) {
-                    return _STD _Find_last_of_pos_vectorized(_Haystack, _Hay_start + 1, _Needle, _Needle_size);
+                if (_Match_try == _Haystack) {
+                    return static_cast<size_t>(-1); // at beginning, no more chance for match
                 }
             }
-#endif // _USE_STD_VECTOR_ALGORITHMS
         }
+
+        // couldn't put one of the characters into the bitmap, fall back to serial algorithm
     }
 
     for (auto _Match_try = _Haystack + _Hay_start;; --_Match_try) {
diff --git a/stl/inc/algorithm b/stl/inc/algorithm
@@ -38,6 +38,15 @@ extern "C" {
 // functions are in native code objects that the compiler cannot analyze. In the absence of the noalias attribute, the
 // compiler has to assume that the denoted arrays are "globally address taken", and that any later calls to
 // unanalyzable routines may modify those arrays.
+const void* __stdcall __std_find_first_of_trivial_1(
+    const void* _First1, const void* _Last1, const void* _First2, const void* _Last2) noexcept;
+const void* __stdcall __std_find_first_of_trivial_2(
+    const void* _First1, const void* _Last1, const void* _First2, const void* _Last2) noexcept;
+const void* __stdcall __std_find_first_of_trivial_4(
+    const void* _First1, const void* _Last1, const void* _First2, const void* _Last2) noexcept;
+const void* __stdcall __std_find_first_of_trivial_8(
+    const void* _First1, const void* _Last1, const void* _First2, const void* _Last2) noexcept;
+
 __declspec(noalias) void __cdecl __std_reverse_copy_trivially_copyable_1(
     const void* _First, const void* _Last, void* _Dest) noexcept;
 __declspec(noalias) void __cdecl __std_reverse_copy_trivially_copyable_2(
@@ -73,6 +82,27 @@ __declspec(noalias) void __stdcall __std_replace_8(
 } // extern "C"
 
 _STD_BEGIN
+template <class _Ty1, class _Ty2>
+_Ty1* _Find_first_of_vectorized(
+    _Ty1* const _First1, _Ty1* const _Last1, _Ty2* const _First2, _Ty2* const _Last2) noexcept {
+    _STL_INTERNAL_STATIC_ASSERT(sizeof(_Ty1) == sizeof(_Ty2));
+    if constexpr (sizeof(_Ty1) == 1) {
+        return const_cast<_Ty1*>(
+            static_cast<const _Ty1*>(::__std_find_first_of_trivial_1(_First1, _Last1, _First2, _Last2)));
+    } else if constexpr (sizeof(_Ty1) == 2) {
+        return const_cast<_Ty1*>(
+            static_cast<const _Ty1*>(::__std_find_first_of_trivial_2(_First1, _Last1, _First2, _Last2)));
+    } else if constexpr (sizeof(_Ty1) == 4) {
+        return const_cast<_Ty1*>(
+            static_cast<const _Ty1*>(::__std_find_first_of_trivial_4(_First1, _Last1, _First2, _Last2)));
+    } else if constexpr (sizeof(_Ty1) == 8) {
+        return const_cast<_Ty1*>(
+            static_cast<const _Ty1*>(::__std_find_first_of_trivial_8(_First1, _Last1, _First2, _Last2)));
+    } else {
+        _STL_INTERNAL_STATIC_ASSERT(false); // unexpected size
+    }
+}
+
 template <size_t _Nx>
 __declspec(noalias) void _Reverse_copy_vectorized(const void* _First, const void* _Last, void* _Dest) noexcept {
     if constexpr (_Nx == 1) {
diff --git a/stl/inc/xutility b/stl/inc/xutility
@@ -98,15 +98,6 @@ const void* __stdcall __std_find_last_trivial_2(const void* _First, const void*
 const void* __stdcall __std_find_last_trivial_4(const void* _First, const void* _Last, uint32_t _Val) noexcept;
 const void* __stdcall __std_find_last_trivial_8(const void* _First, const void* _Last, uint64_t _Val) noexcept;
 
-const void* __stdcall __std_find_first_of_trivial_1(
-    const void* _First1, const void* _Last1, const void* _First2, const void* _Last2) noexcept;
-const void* __stdcall __std_find_first_of_trivial_2(
-    const void* _First1, const void* _Last1, const void* _First2, const void* _Last2) noexcept;
-const void* __stdcall __std_find_first_of_trivial_4(
-    const void* _First1, const void* _Last1, const void* _First2, const void* _Last2) noexcept;
-const void* __stdcall __std_find_first_of_trivial_8(
-    const void* _First1, const void* _Last1, const void* _First2, const void* _Last2) noexcept;
-
 const void* __stdcall __std_search_1(
     const void* _First1, const void* _Last1, const void* _First2, size_t _Count2) noexcept;
 const void* __stdcall __std_search_2(
@@ -252,27 +243,6 @@ _Ty* _Find_last_vectorized(_Ty* const _First, _Ty* const _Last, const _TVal _Val
 // find_first_of vectorization is likely to be a win after this size (in elements)
 _INLINE_VAR constexpr ptrdiff_t _Threshold_find_first_of = 16;
 
-template <class _Ty1, class _Ty2>
-_Ty1* _Find_first_of_vectorized(
-    _Ty1* const _First1, _Ty1* const _Last1, _Ty2* const _First2, _Ty2* const _Last2) noexcept {
-    _STL_INTERNAL_STATIC_ASSERT(sizeof(_Ty1) == sizeof(_Ty2));
-    if constexpr (sizeof(_Ty1) == 1) {
-        return const_cast<_Ty1*>(
-            static_cast<const _Ty1*>(::__std_find_first_of_trivial_1(_First1, _Last1, _First2, _Last2)));
-    } else if constexpr (sizeof(_Ty1) == 2) {
-        return const_cast<_Ty1*>(
-            static_cast<const _Ty1*>(::__std_find_first_of_trivial_2(_First1, _Last1, _First2, _Last2)));
-    } else if constexpr (sizeof(_Ty1) == 4) {
-        return const_cast<_Ty1*>(
-            static_cast<const _Ty1*>(::__std_find_first_of_trivial_4(_First1, _Last1, _First2, _Last2)));
-    } else if constexpr (sizeof(_Ty1) == 8) {
-        return const_cast<_Ty1*>(
-            static_cast<const _Ty1*>(::__std_find_first_of_trivial_8(_First1, _Last1, _First2, _Last2)));
-    } else {
-        _STL_INTERNAL_STATIC_ASSERT(false); // unexpected size
-    }
-}
-
 template <class _Ty1, class _Ty2>
 _Ty1* _Search_vectorized(_Ty1* const _First1, _Ty1* const _Last1, _Ty2* const _First2, const size_t _Count2) noexcept {
     _STL_INTERNAL_STATIC_ASSERT(sizeof(_Ty1) == sizeof(_Ty2));
diff --git a/stl/src/vector_algorithms.cpp b/stl/src/vector_algorithms.cpp
