folly/detail/RangeSse42.cpp

   1 /*
   2  * Copyright 2016 Facebook, Inc.
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  *   http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16
  17 #include "RangeSse42.h"
  18
  19 #include <glog/logging.h>
  20 #include <folly/Portability.h>
  21
  22 //  Essentially, two versions of this file: one with an SSE42 implementation
  23 //  and one with a fallback implementation. We determine which version to use by
  24 //  testing for the presence of the required headers.
  25 //
  26 //  TODO: Maybe this should be done by the build system....
  27 #if !FOLLY_SSE_PREREQ(4, 2)
  28 namespace folly {
  29 namespace detail {
  30 size_t qfind_first_byte_of_sse42(const StringPieceLite haystack,
  31                                  const StringPieceLite needles) {
  32   CHECK(false) << "Function " << __func__ << " only works with SSE42!";
  33   return qfind_first_byte_of_nosse(haystack, needles);
  34 }
  35 }
  36 }
  37 # else
  38 #include <cstdint>
  39 #include <limits>
  40 #include <string>
  41
  42 #include <emmintrin.h>
  43 #include <nmmintrin.h>
  44 #include <smmintrin.h>
  45
  46 #include <folly/Likely.h>
  47
  48 //  GCC 4.9 with ASAN has a problem: a function with no_sanitize_address calling
  49 //  a function with always_inline fails to build. The _mm_* functions are marked
  50 //  always_inline.
  51 //  https://gcc.gnu.org/bugzilla/show_bug.cgi?id=67368
  52 #if defined FOLLY_SANITIZE_ADDRESS && \
  53     FOLLY_SANITIZE_ADDRESS == 1 && \
  54     __GNUC_PREREQ(4, 9)
  55 # define _mm_load_si128(p) (*(p))
  56 # define _mm_loadu_si128(p) ((__m128i)__builtin_ia32_loaddqu((const char*)(p)))
  57 # ifdef _mm_cmpestri
  58 #  undef _mm_cmpestri
  59 # endif
  60 # define _mm_cmpestri(a, b, c, d, e) \
  61   __builtin_ia32_pcmpestri128((__v16qi)(a), b, (__v16qi)(c), d, e)
  62 #endif
  63
  64 namespace folly {
  65 namespace detail {
  66
  67 // It's okay if pages are bigger than this (as powers of two), but they should
  68 // not be smaller.
  69 static constexpr size_t kMinPageSize = 4096;
  70 static_assert(kMinPageSize >= 16,
  71               "kMinPageSize must be at least SSE register size");
  72
  73 template <typename T>
  74 static inline uintptr_t page_for(T* addr) {
  75   return reinterpret_cast<uintptr_t>(addr) / kMinPageSize;
  76 }
  77
  78 static inline size_t nextAlignedIndex(const char* arr) {
  79    auto firstPossible = reinterpret_cast<uintptr_t>(arr) + 1;
  80    return 1 +                       // add 1 because the index starts at 'arr'
  81      ((firstPossible + 15) & ~0xF)  // round up to next multiple of 16
  82      - firstPossible;
  83 }
  84
  85 static size_t qfind_first_byte_of_needles16(const StringPieceLite haystack,
  86                                             const StringPieceLite needles)
  87   FOLLY_DISABLE_ADDRESS_SANITIZER;
  88
  89 // helper method for case where needles.size() <= 16
  90 size_t qfind_first_byte_of_needles16(const StringPieceLite haystack,
  91                                      const StringPieceLite needles) {
  92   DCHECK_GT(haystack.size(), 0);
  93   DCHECK_GT(needles.size(), 0);
  94   DCHECK_LE(needles.size(), 16);
  95   if ((needles.size() <= 2 && haystack.size() >= 256) ||
  96       // must bail if we can't even SSE-load a single segment of haystack
  97       (haystack.size() < 16 &&
  98        page_for(haystack.end() - 1) != page_for(haystack.data() + 15)) ||
  99       // can't load needles into SSE register if it could cross page boundary
 100       page_for(needles.end() - 1) != page_for(needles.data() + 15)) {
 101     return detail::qfind_first_byte_of_nosse(haystack, needles);
 102   }
 103
 104   auto arr2 = _mm_loadu_si128(
 105       reinterpret_cast<const __m128i*>(needles.data()));
 106   // do an unaligned load for first block of haystack
 107   auto arr1 = _mm_loadu_si128(
 108       reinterpret_cast<const __m128i*>(haystack.data()));
 109   auto index = _mm_cmpestri(arr2, needles.size(),
 110                             arr1, haystack.size(), 0);
 111   if (index < 16) {
 112     return index;
 113   }
 114
 115   // Now, we can do aligned loads hereafter...
 116   size_t i = nextAlignedIndex(haystack.data());
 117   for (; i < haystack.size(); i+= 16) {
 118     arr1 =
 119         _mm_load_si128(reinterpret_cast<const __m128i*>(haystack.data() + i));
 120     index = _mm_cmpestri(arr2, needles.size(), arr1, haystack.size() - i, 0);
 121     if (index < 16) {
 122       return i + index;
 123     }
 124   }
 125   return std::string::npos;
 126 }
 127
 128 template <bool HAYSTACK_ALIGNED>
 129 size_t scanHaystackBlock(const StringPieceLite haystack,
 130                          const StringPieceLite needles,
 131                          uint64_t idx)
 132 // Turn off ASAN because the "arr2 = ..." assignment in the loop below reads
 133 // up to 15 bytes beyond end of the buffer in #needles#.  That is ok because
 134 // ptr2 is always 16-byte aligned, so the read can never span a page boundary.
 135 // Also, the extra data that may be read is never actually used.
 136   FOLLY_DISABLE_ADDRESS_SANITIZER;
 137
 138 // Scans a 16-byte block of haystack (starting at blockStartIdx) to find first
 139 // needle. If HAYSTACK_ALIGNED, then haystack must be 16byte aligned.
 140 // If !HAYSTACK_ALIGNED, then caller must ensure that it is safe to load the
 141 // block.
 142 template <bool HAYSTACK_ALIGNED>
 143 size_t scanHaystackBlock(const StringPieceLite haystack,
 144                          const StringPieceLite needles,
 145                          uint64_t blockStartIdx) {
 146   DCHECK_GT(needles.size(), 16);  // should handled by *needles16() method
 147   DCHECK(blockStartIdx + 16 <= haystack.size() ||
 148          (page_for(haystack.data() + blockStartIdx) ==
 149           page_for(haystack.data() + blockStartIdx + 15)));
 150
 151   __m128i arr1;
 152   if (HAYSTACK_ALIGNED) {
 153     arr1 = _mm_load_si128(
 154         reinterpret_cast<const __m128i*>(haystack.data() + blockStartIdx));
 155   } else {
 156     arr1 = _mm_loadu_si128(
 157         reinterpret_cast<const __m128i*>(haystack.data() + blockStartIdx));
 158   }
 159
 160   // This load is safe because needles.size() >= 16
 161   auto arr2 = _mm_loadu_si128(
 162       reinterpret_cast<const __m128i*>(needles.data()));
 163   size_t b = _mm_cmpestri(
 164       arr2, 16, arr1, haystack.size() - blockStartIdx, 0);
 165
 166   size_t j = nextAlignedIndex(needles.data());
 167   for (; j < needles.size(); j += 16) {
 168     arr2 = _mm_load_si128(
 169         reinterpret_cast<const __m128i*>(needles.data() + j));
 170
 171     auto index = _mm_cmpestri(
 172       arr2, needles.size() - j,
 173       arr1, haystack.size() - blockStartIdx, 0);
 174     b = std::min<size_t>(index, b);
 175   }
 176
 177   if (b < 16) {
 178     return blockStartIdx + b;
 179   }
 180   return std::string::npos;
 181 }
 182
 183 size_t qfind_first_byte_of_sse42(const StringPieceLite haystack,
 184                                  const StringPieceLite needles);
 185
 186 size_t qfind_first_byte_of_sse42(const StringPieceLite haystack,
 187                                  const StringPieceLite needles) {
 188   if (UNLIKELY(needles.empty() || haystack.empty())) {
 189     return std::string::npos;
 190   } else if (needles.size() <= 16) {
 191     // we can save some unnecessary load instructions by optimizing for
 192     // the common case of needles.size() <= 16
 193     return qfind_first_byte_of_needles16(haystack, needles);
 194   }
 195
 196   if (haystack.size() < 16 &&
 197       page_for(haystack.end() - 1) != page_for(haystack.data() + 16)) {
 198     // We can't safely SSE-load haystack. Use a different approach.
 199     if (haystack.size() <= 2) {
 200       return qfind_first_byte_of_std(haystack, needles);
 201     }
 202     return qfind_first_byte_of_byteset(haystack, needles);
 203   }
 204
 205   auto ret = scanHaystackBlock<false>(haystack, needles, 0);
 206   if (ret != std::string::npos) {
 207     return ret;
 208   }
 209
 210   size_t i = nextAlignedIndex(haystack.data());
 211   for (; i < haystack.size(); i += 16) {
 212     ret = scanHaystackBlock<true>(haystack, needles, i);
 213     if (ret != std::string::npos) {
 214       return ret;
 215     }
 216   }
 217
 218   return std::string::npos;
 219 }
 220 }
 221 }
 222 #endif