simd::contains (the interfaces) - restored. (facebook#2304)

Summary:
Pull Request resolved: facebook#2304

Faster equivalent to
```
std::ranges::find(haystack, needle) != haystack.end();
```

This diff is just the interfaces.

Differential Revision: D63635498

CMakeLists.txt

@@ -649,11 +649,11 @@ if (BUILD_TESTS OR BUILD_BENCHMARKS)
     DIRECTORY algorithm/simd/detail/test/
       TEST algorithm_simd_detail_simd_any_of_test SOURCES SimdAnyOfTest.cpp
       TEST algorithm_simd_detail_simd_for_each_test SOURCES SimdForEachTest.cpp
+      TEST algorithm_simd_detail_simd_traits_test SOURCES TraitsTest.cpp
       TEST algorithm_simd_detail_unroll_utils_test SOURCES UnrollUtilsTest.cpp
-      # disabled until C++20
-      # TEST algorithm_simd_detail_simd_traits_test SOURCES TraitsTest.cpp
 
     DIRECTORY algorithm/simd/test/
+      TEST algorithm_simd_contains_test SOURCES ContainsTest.cpp
       TEST algorithm_simd_find_fixed_test SOURCES FindFixedTest.cpp
       TEST algorithm_simd_movemask_test SOURCES MovemaskTest.cpp
 

folly/algorithm/simd/BUCK

@@ -22,3 +22,16 @@ cpp_library(
         "//folly/algorithm/simd/detail:traits",
     ],
 )
+
+cpp_library(
+    name = "contains",
+    srcs = ["Contains.cpp"],
+    headers = ["Contains.h"],
+    deps = [
+        "//folly/algorithm/simd/detail:simd_contains_impl",
+    ],
+    exported_deps = [
+        "//folly:c_portability",
+        "//folly/algorithm/simd/detail:traits",
+    ],
+)

folly/algorithm/simd/Contains.cpp

@@ -0,0 +1,42 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * 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 <folly/algorithm/simd/Contains.h>
+
+#include <algorithm>
+#include <cstring>
+#include <folly/algorithm/simd/detail/ContainsImpl.h>
+
+namespace folly::simd::detail {
+
+bool containsU8(folly::span<const std::uint8_t> haystack, std::uint8_t needle) {
+  return containsImpl(haystack, needle);
+}
+bool containsU16(
+    folly::span<const std::uint16_t> haystack, std::uint16_t needle) {
+  return containsImpl(haystack, needle);
+}
+bool containsU32(
+    folly::span<const std::uint32_t> haystack, std::uint32_t needle) {
+  return containsImpl(haystack, needle);
+}
+
+bool containsU64(
+    folly::span<const std::uint64_t> haystack, std::uint64_t needle) {
+  return containsImpl(haystack, needle);
+}
+
+} // namespace folly::simd::detail

folly/algorithm/simd/Contains.h

@@ -0,0 +1,80 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * 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 <folly/CPortability.h>
+#include <folly/algorithm/simd/detail/Traits.h>
+
+#include <iterator>
+
+namespace folly::simd {
+namespace detail {
+
+// no overloading for easier profiling.
+
+bool containsU8(folly::span<const std::uint8_t> haystack, std::uint8_t needle);
+bool containsU16(
+    folly::span<const std::uint16_t> haystack, std::uint16_t needle);
+bool containsU32(
+    folly::span<const std::uint32_t> haystack, std::uint32_t needle);
+bool containsU64(
+    folly::span<const std::uint64_t> haystack, std::uint64_t needle);
+
+template <typename R>
+using std_range_value_t = typename std::iterator_traits<decltype(std::begin(
+    std::declval<R&>()))>::value_type;
+
+} // namespace detail
+
+/**
+ * folly::simd::contains
+ * folly::simd::contains_fn
+ *
+ * A vectorized version of `std::ranges::find(r, x) != r.end()`.
+ * Only works for "simd friendly cases" -
+ * specifically the ones where the type can be reasonably cast
+ * to `uint8/16/32/64_t`.
+ *
+ **/
+struct contains_fn {
+  template <
+      typename R,
+      typename = std::enable_if_t<
+          std::is_invocable_v<detail::AsSimdFriendlyUintFn, R>>>
+  FOLLY_ERASE bool operator()(R&& r, detail::std_range_value_t<R> x) const {
+    auto castR = detail::asSimdFriendlyUint(folly::span(r));
+    auto castX = detail::asSimdFriendlyUint(x);
+
+    using T = decltype(castX);
+
+    if constexpr (std::is_same_v<T, std::uint8_t>) {
+      return detail::containsU8(castR, castX);
+    } else if constexpr (std::is_same_v<T, std::uint16_t>) {
+      return detail::containsU16(castR, castX);
+    } else if constexpr (std::is_same_v<T, std::uint32_t>) {
+      return detail::containsU32(castR, castX);
+    } else {
+      static_assert(
+          std::is_same_v<T, std::uint64_t>, "internal error, unknown type");
+      return detail::containsU64(castR, castX);
+    }
+  }
+};
+
+inline constexpr contains_fn contains;
+
+} // namespace folly::simd

folly/algorithm/simd/detail/BUCK

@@ -26,6 +26,17 @@ cpp_library(
     ],
 )
 
+cpp_library(
+    name = "simd_contains_impl",
+    headers = ["ContainsImpl.h"],
+    exported_deps = [
+        ":simd_any_of",
+        ":simd_char_platform",
+        "//folly:c_portability",
+        "//folly/container:span",
+    ],
+)
+
 cpp_library(
     name = "simd_for_each",
     headers = ["SimdForEach.h"],

folly/algorithm/simd/detail/ContainsImpl.h

@@ -0,0 +1,91 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * 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 <algorithm>
+#include <cstring>
+#include <cwchar>
+#include <type_traits>
+
+#include <folly/CPortability.h>
+#include <folly/algorithm/simd/detail/SimdAnyOf.h>
+#include <folly/algorithm/simd/detail/SimdCharPlatform.h>
+#include <folly/container/span.h>
+
+namespace folly::simd::detail {
+
+/*
+ * The functions in this file are FOLLY_ERASE to make sure
+ * that the only place behind a call boundary is the explicit one.
+ */
+
+template <typename T>
+FOLLY_ERASE bool containsImplStd(folly::span<const T> haystack, T needle) {
+  static_assert(
+      std::is_unsigned_v<T>, "we should only get here for uint8/16/32/64");
+  if constexpr (sizeof(T) == 1) {
+    auto* ptr = reinterpret_cast<const char*>(haystack.data());
+    auto castNeedle = static_cast<char>(needle);
+    if (haystack.empty()) { // memchr requires not null
+      return false;
+    }
+    return std::memchr(ptr, castNeedle, haystack.size()) != nullptr;
+  } else if constexpr (sizeof(T) == sizeof(wchar_t)) {
+    auto* ptr = reinterpret_cast<const wchar_t*>(haystack.data());
+    auto castNeedle = static_cast<wchar_t>(needle);
+    if (haystack.empty()) { // wmemchr requires not null
+      return false;
+    }
+    return std::wmemchr(ptr, castNeedle, haystack.size()) != nullptr;
+  } else {
+    // Using find instead of any_of on an off chance that the standard library
+    // will add some custom vectorization.
+    // That wouldn't be possible for any_of because of the predicates.
+    return std::find(haystack.begin(), haystack.end(), needle) !=
+        haystack.end();
+  }
+}
+
+template <typename T>
+constexpr bool hasHandwrittenContains() {
+  return std::is_same_v<T, std::uint8_t> &&
+      !std::is_same_v<SimdCharPlatform, void>;
+}
+
+template <typename T>
+FOLLY_ERASE bool containsImplHandwritten(
+    folly::span<const T> haystack, T needle) {
+  static_assert(std::is_same_v<T, std::uint8_t>, "");
+  auto as_chars = folly::reinterpret_span_cast<const char>(haystack);
+  return simdAnyOf<SimdCharPlatform, 4>(
+      as_chars.data(),
+      as_chars.data() + as_chars.size(),
+      [&](SimdCharPlatform::reg_t x) {
+        return SimdCharPlatform::equal(x, static_cast<char>(needle));
+      });
+}
+
+template <typename T>
+FOLLY_ERASE bool containsImpl(folly::span<const T> haystack, T needle) {
+  if constexpr (hasHandwrittenContains<T>()) {
+    return containsImplHandwritten(haystack, needle);
+  } else {
+    return containsImplStd(haystack, needle);
+  }
+}
+
+} // namespace folly::simd::detail

folly/algorithm/simd/detail/test/TraitsTest.cpp

@@ -24,7 +24,7 @@
 
 namespace folly::simd::detail {
 
-struct FollySimdTraitsTest : testing::Test {};
+struct SimdTraitsTest : testing::Test {};
 
 namespace simd_friendly_equivalent_scalar_test {
 
@@ -161,7 +161,7 @@ static_assert(!std::is_invocable_v<AsSimdFriendlyUintFn, std::set<int>>);
 
 } // namespace as_simd_friendly_uint_type_test
 
-TEST_F(FollySimdTraitsTest, AsSimdFriendly) {
+TEST_F(SimdTraitsTest, AsSimdFriendly) {
   enum SomeEnum : int { Foo = 1, Bar, Baz };
 
   static_assert(asSimdFriendly(SomeEnum::Foo) == 1);
@@ -171,7 +171,7 @@ TEST_F(FollySimdTraitsTest, AsSimdFriendly) {
   ASSERT_THAT(castSpan, testing::ElementsAre(1, 2, 3));
 }
 
-TEST_F(FollySimdTraitsTest, AsSimdFriendlyUint) {
+TEST_F(SimdTraitsTest, AsSimdFriendlyUint) {
   enum SomeEnum : int { Foo = 1, Bar, Baz };
 
   static_assert(asSimdFriendlyUint(SomeEnum::Foo) == 1U);

folly/algorithm/simd/test/BUCK

@@ -36,3 +36,14 @@ cpp_benchmark(
         "//folly/init:init",
     ],
 )
+
+cpp_unittest(
+    name = "contains_test",
+    srcs = ["ContainsTest.cpp"],
+    headers = [],
+    deps = [
+        "//folly/algorithm/simd:contains",
+        "//folly/algorithm/simd/detail:simd_contains_impl",
+        "//folly/portability:gtest",
+    ],
+)