simd::contains (the interfaces) (#2299)

Summary:
Pull Request resolved: #2299

simdContains - everything but the actual handwritten algorithm.

Reviewed By: Gownta

Differential Revision: D63116101

fbshipit-source-id: 2c9b23f0111f0fa2f703ca13e8cd3a1097c685fd

CMakeLists.txt

@@ -649,13 +649,13 @@ 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_find_fixed_test SOURCES FindFixedTest.cpp
       TEST algorithm_simd_movemask_test SOURCES MovemaskTest.cpp
+      TEST algorithm_simd_simd_contains_test SOURCES SimdContainsTest.cpp
 
     DIRECTORY chrono/test/
       TEST chrono_conv_test WINDOWS_DISABLED

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,65 @@
+/*
+ * 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 <ranges>
+
+namespace folly::simd {
+namespace detail {
+
+// no overloading for easier of 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);
+
+} // namespace detail
+
+struct contains_fn {
+  template <std::ranges::contiguous_range R>
+    requires simd::detail::has_integral_simd_friendly_equivalent<
+        std::ranges::range_value_t<R>>
+  FOLLY_ERASE bool operator()(R&& rng, std::ranges::range_value_t<R> x) const {
+    auto castRng = simd::detail::asSimdFriendlyUint(folly::span(rng));
+    auto castX = simd::detail::asSimdFriendlyUint(x);
+
+    using T = decltype(castX);
+
+    if constexpr (std::is_same_v<T, std::uint8_t>) {
+      return detail::containsU8(castRng, castX);
+    } else if constexpr (std::is_same_v<T, std::uint16_t>) {
+      return detail::containsU16(castRng, castX);
+    } else if constexpr (std::is_same_v<T, std::uint32_t>) {
+      return detail::containsU32(castRng, castX);
+    } else {
+      static_assert(
+          std::is_same_v<T, std::uint64_t>, "internal error, unknown type");
+      return detail::containsU64(castRng, castX);
+    }
+  }
+};
+
+inline constexpr contains_fn contains;
+
+} // namespace folly::simd

folly/algorithm/simd/FindFixed.h

@@ -293,7 +293,8 @@ constexpr std::optional<std::size_t> findFixed(std::span<const T, N> where, U x)
     return find_fixed_detail::findFixedConstexpr(std::span<const T>(where), x);
   } else {
     return find_fixed_detail::findFixedDispatch(
-        detail::asSimdFriendlyUint(where), detail::asSimdFriendlyUint(x));
+        simd::detail::asSimdFriendlyUint(where),
+        simd::detail::asSimdFriendlyUint(x));
   }
 }
 

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"],
@@ -40,6 +51,7 @@ cpp_library(
     name = "traits",
     headers = ["Traits.h"],
     exported_deps = [
+        "//folly:c_portability",
         "//folly:memory",
         "//folly:traits",
         "//folly/container:span",

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/Traits.h

@@ -16,14 +16,15 @@
 
 #pragma once
 
+#include <folly/CPortability.h>
 #include <folly/Memory.h>
 #include <folly/Traits.h>
 #include <folly/container/span.h>
 
 #include <concepts>
 #include <type_traits>
 
-namespace folly::detail {
+namespace folly::simd::detail {
 
 template <typename T>
 auto findSimdFriendlyEquivalent() {
@@ -36,25 +37,9 @@ auto findSimdFriendlyEquivalent() {
       return double{};
     }
   } else if constexpr (std::is_signed_v<T>) {
-    if constexpr (sizeof(T) == 1) {
-      return std::int8_t{};
-    } else if constexpr (sizeof(T) == 2) {
-      return std::int16_t{};
-    } else if constexpr (sizeof(T) == 4) {
-      return std::int32_t{};
-    } else if constexpr (sizeof(T) == 8) {
-      return std::int64_t{};
-    }
+    return int_bits_t<sizeof(T) * 8>{};
   } else if constexpr (std::is_unsigned_v<T>) {
-    if constexpr (sizeof(T) == 1) {
-      return std::uint8_t{};
-    } else if constexpr (sizeof(T) == 2) {
-      return std::uint16_t{};
-    } else if constexpr (sizeof(T) == 4) {
-      return std::uint32_t{};
-    } else if constexpr (sizeof(T) == 8) {
-      return std::uint64_t{};
-    }
+    return uint_bits_t<sizeof(T) * 8>{};
   }
 }
 
@@ -63,7 +48,7 @@ concept has_simd_friendly_equivalent =
     !std::is_same_v<void, decltype(findSimdFriendlyEquivalent<T>())>;
 
 template <has_simd_friendly_equivalent T>
-using simd_friendly_equivalent_t = folly::like_t< //
+using simd_friendly_equivalent_t = like_t< //
     T,
     decltype(findSimdFriendlyEquivalent<std::remove_const_t<T>>())>;
 
@@ -77,24 +62,23 @@ template <has_integral_simd_friendly_equivalent T>
 using integral_simd_friendly_equivalent = simd_friendly_equivalent_t<T>;
 
 template <has_simd_friendly_equivalent T, std::size_t Extend>
-auto asSimdFriendly(folly::span<T, Extend> s) {
-  return folly::reinterpret_span_cast<simd_friendly_equivalent_t<T>>(s);
+FOLLY_ERASE auto asSimdFriendly(folly::span<T, Extend> s) {
+  return reinterpret_span_cast<simd_friendly_equivalent_t<T>>(s);
 }
 
 template <has_simd_friendly_equivalent T>
-constexpr auto asSimdFriendly(T x) {
+FOLLY_ERASE constexpr auto asSimdFriendly(T x) {
   return static_cast<simd_friendly_equivalent_t<T>>(x);
 }
 
 template <has_simd_friendly_equivalent T, std::size_t Extend>
-auto asSimdFriendlyUint(folly::span<T, Extend> s) {
-  return folly::reinterpret_span_cast<
-      folly::like_t<T, uint_bits_t<sizeof(T) * 8>>>(s);
+FOLLY_ERASE auto asSimdFriendlyUint(folly::span<T, Extend> s) {
+  return reinterpret_span_cast<like_t<T, uint_bits_t<sizeof(T) * 8>>>(s);
 }
 
 template <has_simd_friendly_equivalent T>
-constexpr auto asSimdFriendlyUint(T x) {
+FOLLY_ERASE constexpr auto asSimdFriendlyUint(T x) {
   return static_cast<uint_bits_t<sizeof(T) * 8>>(x);
 }
 
-} // namespace folly::detail
+} // namespace folly::simd::detail

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

@@ -19,9 +19,9 @@
 #include <folly/portability/GMock.h>
 #include <folly/portability/GTest.h>
 
-namespace folly::detail {
+namespace folly::simd::detail {
 
-struct FollySimdTraitsTest : testing::Test {};
+struct SimdTraitsTest : testing::Test {};
 
 namespace simd_friendly_equivalent_test {
 
@@ -98,7 +98,7 @@ static_assert(Overloading{}(SomeInt{}) == 2);
 
 } // namespace integral_simd_friendly_equivalent_test
 
-TEST_F(FollySimdTraitsTest, AsSimdFriendly) {
+TEST_F(SimdTraitsTest, AsSimdFriendly) {
   enum SomeEnum : int { Foo = 1, Bar, Baz };
 
   static_assert(asSimdFriendly(SomeEnum::Foo) == 1);
@@ -122,7 +122,7 @@ void asSimdFriendlyUintTypeTest() {
       asSimdFriendlyUint(std::span<const From>{}), std::span<const To>{});
 }
 
-TEST_F(FollySimdTraitsTest, AsSimdFriendlyUint) {
+TEST_F(SimdTraitsTest, AsSimdFriendlyUint) {
   enum SomeEnum : int { Foo = 1, Bar, Baz };
 
   static_assert(asSimdFriendlyUint(SomeEnum::Foo) == 1U);
@@ -136,4 +136,4 @@ TEST_F(FollySimdTraitsTest, AsSimdFriendlyUint) {
   asSimdFriendlyUintTypeTest<double, std::uint64_t>();
 }
 
-} // namespace folly::detail
+} // namespace folly::simd::detail