Optimize mem utils functions

- add `avx2_strstr` to accelerate substr search
- add `avx2_mem_equal` to accelerate mem equal cmp

dbms/src/Functions/CMakeLists.txt

@@ -19,6 +19,10 @@ add_headers_and_sources(clickhouse_functions ./GatherUtils)
 add_headers_and_sources(clickhouse_functions ./Conditional)
 add_headers_and_sources(clickhouse_functions ${TiFlash_BINARY_DIR}/dbms/src/Functions)
 
+if (TIFLASH_ENABLE_AVX_SUPPORT)
+    set_source_files_properties(CollationStringOptimized.cpp APPEND COMPILE_FLAGS "-mavx -mavx2")
+endif ()
+
 list(REMOVE_ITEM clickhouse_functions_sources IFunction.cpp FunctionFactory.cpp FunctionHelpers.cpp)
 list(REMOVE_ITEM clickhouse_functions_headers IFunction.h FunctionFactory.h FunctionHelpers.h)
 

dbms/src/Functions/CollationStringOptimized.cpp

@@ -0,0 +1,53 @@
+// Copyright 2022 PingCAP, Ltd.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <Functions/CollationStringSearch.h>
+#include <Functions/CollationStringSearchOptimized.h>
+
+#include <cstddef>
+#include <cstdint>
+#include <limits>
+
+namespace DB
+{
+
+template <bool revert>
+bool StringPatternMatch(
+    const ColumnString::Chars_t & a_data,
+    const ColumnString::Offsets & a_offsets,
+    const std::string_view & pattern_str,
+    uint8_t escape_char,
+    const TiDB::TiDBCollatorPtr & collator,
+    PaddedPODArray<UInt8> & c)
+{
+    return StringPatternMatchImpl<revert>(a_data, a_offsets, pattern_str, escape_char, collator, c);
+}
+
+template bool StringPatternMatch<true>(
+    const ColumnString::Chars_t & a_data,
+    const ColumnString::Offsets & a_offsets,
+    const std::string_view & pattern_str,
+    uint8_t escape_char,
+    const TiDB::TiDBCollatorPtr & collator,
+    PaddedPODArray<UInt8> & c);
+
+template bool StringPatternMatch<false>(
+    const ColumnString::Chars_t & a_data,
+    const ColumnString::Offsets & a_offsets,
+    const std::string_view & pattern_str,
+    uint8_t escape_char,
+    const TiDB::TiDBCollatorPtr & collator,
+    PaddedPODArray<UInt8> & c);
+
+} // namespace DB

dbms/src/Functions/CollationStringSearch.h

@@ -0,0 +1,31 @@
+// Copyright 2022 PingCAP, Ltd.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <Columns/ColumnString.h>
+
+namespace DB
+{
+
+template <bool revert>
+extern bool StringPatternMatch(
+    const ColumnString::Chars_t & a_data,
+    const ColumnString::Offsets & a_offsets,
+    const std::string_view & pattern_str,
+    uint8_t escape_char,
+    const TiDB::TiDBCollatorPtr & collator,
+    PaddedPODArray<UInt8> & c);
+
+} // namespace DB

dbms/src/Functions/CollationStringSearchOptimized.h

@@ -14,11 +14,10 @@
 
 #pragma once
 
-#include <Common/StringSearcher.h>
 #include <Common/UTF8Helpers.h>
 #include <Functions/CollationOperatorOptimized.h>
 #include <Storages/Transaction/CollatorUtils.h>
-#include <sys/types.h>
+#include <common/mem_utils_opt.h>
 
 #include <algorithm>
 #include <cassert>
@@ -31,6 +30,19 @@ namespace TiDB
 static constexpr char ANY = '%';
 static constexpr char ONE = '_';
 
+FLATTEN_INLINE_PURE static inline size_t BuiltinStrFind(std::string_view src, std::string_view needle)
+{
+#ifdef TIFLASH_ENABLE_AVX_SUPPORT
+#ifdef TIFLASH_USE_AVX2_COMPILE_FLAG
+    return mem_utils::avx2_strstr(src, needle);
+#else
+    return avx2_strstr(src, needle);
+#endif
+#else
+    return src.find(needle);
+#endif
+}
+
 /*
     Unicode Code    UTF-8 Code
     0000～007F      0xxxxxxx
@@ -88,7 +100,7 @@ struct BinStrPattern
 
         auto last_match_start = std::string_view::npos;
 
-        const auto & fn_try_add_last_match_str = [&](size_t end_offset) {
+        auto && fn_try_add_last_match_str = [&](size_t end_offset) {
             if (last_match_start != std::string_view::npos)
             {
                 match_sub_str.emplace_back(&pattern[last_match_start], end_offset - last_match_start);
@@ -194,15 +206,15 @@ struct BinStrPattern
 
     // check str equality
     // - make src invalid if remain size if smaller than required
-    bool matchStrEqual(const std::string_view & src, MatchDesc & desc) const
+    ALWAYS_INLINE inline bool matchStrEqual(const std::string_view & src, MatchDesc & desc) const
     {
         const auto & match_str = match_sub_str[desc.match_str_index_start];
         if (desc.srcSize() < match_str.size())
         {
             desc.makeSrcInvalid();
             return false;
         }
-        if (DB::RawStrEqualCompare(desc.getSrcStrView(src.data(), match_str.size()), match_str))
+        if (!DB::IsRawStrEqual(desc.getSrcStrView(src.data(), match_str.size()), match_str))
         {
             return false;
         }
@@ -214,7 +226,7 @@ struct BinStrPattern
     // match from start exactly
     // - return true if meet %
     // - return false if failed to match else true
-    bool matchExactly(const std::string_view & src, MatchDesc & cur_match_desc) const
+    ALWAYS_INLINE inline bool matchExactly(const std::string_view & src, MatchDesc & cur_match_desc) const
     {
         // match from start
         for (; !cur_match_desc.patternEmpty(); cur_match_desc.pattern_index_start++)
@@ -245,7 +257,7 @@ struct BinStrPattern
     // match from end exactly
     // - return true if meet %
     // - return false if failed to match else true
-    bool matchExactlyReverse(const std::string_view & src, MatchDesc & cur_match_desc) const
+    ALWAYS_INLINE inline bool matchExactlyReverse(const std::string_view & src, MatchDesc & cur_match_desc) const
     {
         for (; !cur_match_desc.patternEmpty(); --cur_match_desc.pattern_index_end)
         {
@@ -263,7 +275,7 @@ struct BinStrPattern
                     return false;
                 }
 
-                if (DB::RawStrEqualCompare({src.data() + cur_match_desc.src_index_end - match_str.size(), match_str.size()}, match_str))
+                if (!DB::IsRawStrEqual({src.data() + cur_match_desc.src_index_end - match_str.size(), match_str.size()}, match_str))
                 {
                     return false;
                 }
@@ -286,7 +298,7 @@ struct BinStrPattern
     // search by pattern `...%..%`
     // - return true if meet %
     // - return false if failed to search
-    bool searchByPattern(const std::string_view & src, MatchDesc & desc) const
+    ALWAYS_INLINE inline bool searchByPattern(const std::string_view & src, MatchDesc & desc) const
     {
         assert(match_types[desc.pattern_index_end - 1] == MatchType::Any);
         assert(!desc.patternEmpty());
@@ -320,10 +332,8 @@ struct BinStrPattern
 
             // search sub str
             // - seachers like `ASCIICaseSensitiveStringSearcher` or `Volnitsky` are too heavy for small str
-            // - TODO: optimize strstr search by simd
             {
-                pos = src_view.find(match_str);
-                // pos = sse2_strstr(src_view, match_str);
+                pos = BuiltinStrFind(src_view, match_str);
             }
 
             if (pos == std::string_view::npos)
@@ -356,7 +366,7 @@ struct BinStrPattern
         }
     };
 
-    bool match(std::string_view src) const
+    ALWAYS_INLINE inline bool match(std::string_view src) const
     {
         MatchDesc cur_match_desc;
         {
@@ -421,7 +431,7 @@ struct BinStrPattern
 namespace DB
 {
 template <typename Result, bool revert, bool utf8>
-ALWAYS_INLINE inline void BinStringPatternMatch(
+inline void BinStringPatternMatch(
     const ColumnString::Chars_t & a_data,
     const ColumnString::Offsets & a_offsets,
     const std::string_view & pattern_str,
@@ -436,7 +446,7 @@ ALWAYS_INLINE inline void BinStringPatternMatch(
 }
 
 template <bool revert, typename Result>
-ALWAYS_INLINE inline bool StringPatternMatch(
+ALWAYS_INLINE inline bool StringPatternMatchImpl(
     const ColumnString::Chars_t & a_data,
     const ColumnString::Offsets & a_offsets,
     const std::string_view & pattern_str,

dbms/src/Functions/FunctionsStringSearch.cpp

@@ -17,7 +17,8 @@
 #include <Common/Volnitsky.h>
 #include <Common/config.h>
 #include <DataTypes/DataTypeFixedString.h>
-#include <Functions/CollationStringSearchOptimized.h>
+#include <Functions/CollationOperatorOptimized.h>
+#include <Functions/CollationStringSearch.h>
 #include <Functions/FunctionFactory.h>
 #include <Functions/FunctionsStringSearch.h>
 #include <Functions/Regexps.h>

dbms/src/Storages/Transaction/CollatorUtils.h

@@ -16,7 +16,9 @@
 
 #include <common/StringRef.h>
 #include <common/defines.h>
+#include <common/mem_utils_opt.h>
 
+#include <cstring>
 #include <memory>
 
 namespace DB
@@ -28,12 +30,28 @@ ALWAYS_INLINE inline int signum(T val)
     return (0 < val) - (val < 0);
 }
 
+FLATTEN_INLINE_PURE static inline bool IsRawStrEqual(const std::string_view & lhs, const std::string_view & rhs)
+{
+    if (lhs.size() != rhs.size())
+        return false;
+
+#ifdef TIFLASH_ENABLE_AVX_SUPPORT
+#ifdef TIFLASH_USE_AVX2_COMPILE_FLAG
+    return mem_utils::avx2_mem_equal(lhs.data(), rhs.data(), lhs.size());
+#else
+    return avx2_mem_equal(lhs.data(), rhs.data(), lhs.size());
+#endif
+#else
+    return 0 == std::memcmp(lhs.data(), rhs.data(), lhs.size());
+#endif
+}
+
 // Check equality is much faster than other comparison.
 // - check size first
 // - return 0 if equal else 1
-FLATTEN_INLINE_PURE inline int RawStrEqualCompare(const std::string_view & lhs, const std::string_view & rhs)
+FLATTEN_INLINE_PURE static inline int RawStrEqualCompare(const std::string_view & lhs, const std::string_view & rhs)
 {
-    return StringRef(lhs) == StringRef(rhs) ? 0 : 1;
+    return IsRawStrEqual(lhs, rhs) ? 0 : 1;
 }
 
 // Compare str view by memcmp

dbms/src/Storages/Transaction/tests/gtest_tidb_collator.cpp

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <Functions/CollationStringSearch.h>
 #include <Functions/CollationStringSearchOptimized.h>
 #include <Storages/Transaction/Collator.h>
 #include <Storages/Transaction/CollatorUtils.h>
@@ -225,7 +226,7 @@ void testCollator()
 
             ColumnString::Chars_t strs;
             ColumnString::Offsets offsets;
-            std::vector<bool> res;
+            PaddedPODArray<UInt8> res;
             { // init data
                 ColumnString::Offset current_new_offset = 0;
                 for (const auto & inner_c : inner_cases)

libs/libcommon/CMakeLists.txt

@@ -52,6 +52,7 @@ add_library (common ${SPLIT_SHARED}
     src/crc64.cpp
     src/simd.cpp
     src/detect_features.cpp
+    src/avx2_mem_utils_impl.cpp
 
     include/common/types.h
     include/common/DateLUT.h
@@ -159,6 +160,7 @@ endif ()
 if (TIFLASH_ENABLE_AVX_SUPPORT)
     # https://gcc.gnu.org/onlinedocs/gcc/x86-Options.html
     set_source_files_properties(src/mem_utils_avx2.cpp APPEND COMPILE_FLAGS "-mavx -mavx2")
+    set_source_files_properties(src/avx2_mem_utils_impl.cpp APPEND COMPILE_FLAGS "-mavx -mavx2")
     if (TIFLASH_COMPILER_VPCLMULQDQ_SUPPORT)
         set_source_files_properties(src/crc64_avx2.cpp
                 APPEND COMPILE_FLAGS "-mavx2 -mpclmul -mvpclmulqdq -Wno-ignored-attributes")