From c48964c77c27f9ea5eb38357c82e05ebc68a1b88 Mon Sep 17 00:00:00 2001
From: Denis Yaroshevskiy <dyaroshev@meta.com>
Date: Mon, 23 Sep 2024 06:51:52 -0700
Subject: [PATCH] simdContains (the interfaces) (#2299)

Summary:
Pull Request resolved: https://github.com/facebook/folly/pull/2299

simdContains - everything but the actual handwritten algorithm.

Differential Revision: D63116101
---
 CMakeLists.txt                                |  1 +
 folly/algorithm/simd/BUCK                     | 12 +++
 folly/algorithm/simd/SimdContains.cpp         | 43 +++++++++
 folly/algorithm/simd/SimdContains.h           | 63 ++++++++++++++
 folly/algorithm/simd/detail/BUCK              | 11 +++
 .../algorithm/simd/detail/SimdContainsImpl.h  | 87 +++++++++++++++++++
 folly/algorithm/simd/test/BUCK                | 11 +++
 .../algorithm/simd/test/SimdContainsTest.cpp  | 78 +++++++++++++++++
 8 files changed, 306 insertions(+)
 create mode 100644 folly/algorithm/simd/SimdContains.cpp
 create mode 100644 folly/algorithm/simd/SimdContains.h
 create mode 100644 folly/algorithm/simd/detail/SimdContainsImpl.h
 create mode 100644 folly/algorithm/simd/test/SimdContainsTest.cpp

diff --git a/CMakeLists.txt b/CMakeLists.txt
index c30fc9938ed..10a6d13acb5 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -652,6 +652,7 @@ if (BUILD_TESTS OR BUILD_BENCHMARKS)
       TEST algorithm_simd_detail_unroll_utils_test SOURCES UnrollUtilsTest.cpp
       # disabled until C++20
       # TEST algorithm_simd_detail_simd_traits_test SOURCES TraitsTest.cpp
+      # TEST algorithm_simd_detail_simd_contains_test SOURCES SimdContainsTest.cpp
 
     DIRECTORY algorithm/simd/test/
       TEST algorithm_simd_find_fixed_test SOURCES FindFixedTest.cpp
diff --git a/folly/algorithm/simd/BUCK b/folly/algorithm/simd/BUCK
index 1343e8f9b73..9dd7881040d 100644
--- a/folly/algorithm/simd/BUCK
+++ b/folly/algorithm/simd/BUCK
@@ -22,3 +22,15 @@ cpp_library(
         "//folly/algorithm/simd/detail:traits",
     ],
 )
+
+cpp_library(
+    name = "simd_contains",
+    srcs = ["SimdContains.cpp"],
+    headers = ["SimdContains.h"],
+    deps = [
+        "//folly/algorithm/simd/detail:simd_contains_impl",
+    ],
+    exported_deps = [
+        "//folly/algorithm/simd/detail:traits",
+    ],
+)
diff --git a/folly/algorithm/simd/SimdContains.cpp b/folly/algorithm/simd/SimdContains.cpp
new file mode 100644
index 00000000000..554d74dccd1
--- /dev/null
+++ b/folly/algorithm/simd/SimdContains.cpp
@@ -0,0 +1,43 @@
+/*
+ * 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/SimdContains.h>
+
+#include <algorithm>
+#include <cstring>
+#include <folly/algorithm/simd/detail/SimdContainsImpl.h>
+
+namespace folly::simd_detail {
+
+bool simdContainsU8(
+    folly::span<const std::uint8_t> haystack, std::uint8_t needle) {
+  return simdContainsImpl(haystack, needle);
+}
+bool simdContainsU16(
+    folly::span<const std::uint16_t> haystack, std::uint16_t needle) {
+  return simdContainsImpl(haystack, needle);
+}
+bool simdContainsU32(
+    folly::span<const std::uint32_t> haystack, std::uint32_t needle) {
+  return simdContainsImpl(haystack, needle);
+}
+
+bool simdContainsU64(
+    folly::span<const std::uint64_t> haystack, std::uint64_t needle) {
+  return simdContainsImpl(haystack, needle);
+}
+
+} // namespace folly::simd_detail
diff --git a/folly/algorithm/simd/SimdContains.h b/folly/algorithm/simd/SimdContains.h
new file mode 100644
index 00000000000..900999adaa7
--- /dev/null
+++ b/folly/algorithm/simd/SimdContains.h
@@ -0,0 +1,63 @@
+/*
+ * 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/algorithm/simd/detail/Traits.h>
+
+#include <ranges>
+
+namespace folly {
+namespace simd_detail {
+
+// no overloading for easier of profiling.
+
+bool simdContainsU8(
+    folly::span<const std::uint8_t> haystack, std::uint8_t needle);
+bool simdContainsU16(
+    folly::span<const std::uint16_t> haystack, std::uint16_t needle);
+bool simdContainsU32(
+    folly::span<const std::uint32_t> haystack, std::uint32_t needle);
+bool simdContainsU64(
+    folly::span<const std::uint64_t> haystack, std::uint64_t needle);
+
+} // namespace simd_detail
+
+struct simd_contains_fn {
+  template <std::ranges::contiguous_range R>
+    requires detail::has_integral_simd_friendly_equivalent<
+        std::ranges::range_value_t<R>>
+  bool operator()(R&& rng, std::ranges::range_value_t<R> x) const {
+    auto castRng = detail::asSimdFriendlyUint(folly::span(rng));
+    auto castX = detail::asSimdFriendlyUint(x);
+
+    using T = decltype(castX);
+
+    if constexpr (std::is_same_v<T, std::uint8_t>) {
+      return simd_detail::simdContainsU8(castRng, castX);
+    } else if constexpr (std::is_same_v<T, std::uint16_t>) {
+      return simd_detail::simdContainsU16(castRng, castX);
+    } else if constexpr (std::is_same_v<T, std::uint32_t>) {
+      return simd_detail::simdContainsU32(castRng, castX);
+    } else {
+      static_assert(
+          std::is_same_v<T, std::uint64_t>, "internal error, unknown type");
+      return simd_detail::simdContainsU64(castRng, castX);
+    }
+  }
+} inline constexpr simd_contains;
+
+} // namespace folly
diff --git a/folly/algorithm/simd/detail/BUCK b/folly/algorithm/simd/detail/BUCK
index 53a172a2e6b..ef44c599b11 100644
--- a/folly/algorithm/simd/detail/BUCK
+++ b/folly/algorithm/simd/detail/BUCK
@@ -26,6 +26,17 @@ cpp_library(
     ],
 )
 
+cpp_library(
+    name = "simd_contains_impl",
+    headers = ["SimdContainsImpl.h"],
+    exported_deps = [
+        ":simd_any_of",
+        ":simd_char_platform",
+        "//folly:c_portability",
+        "//folly/container:span",
+    ],
+)
+
 cpp_library(
     name = "simd_for_each",
     headers = ["SimdForEach.h"],
diff --git a/folly/algorithm/simd/detail/SimdContainsImpl.h b/folly/algorithm/simd/detail/SimdContainsImpl.h
new file mode 100644
index 00000000000..ff0f0502134
--- /dev/null
+++ b/folly/algorithm/simd/detail/SimdContainsImpl.h
@@ -0,0 +1,87 @@
+/*
+ * 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 funcitons in this file are FOLLY_ALWAYS_INLINE to make sure
+ * that the only place behind a call boundary is the explicit one.
+ */
+
+template <typename T>
+FOLLY_ALWAYS_INLINE bool simdContainsImplStd(
+    folly::span<const T> haystack, T needle) {
+  if constexpr (sizeof(T) == 1) {
+    auto* ptr = reinterpret_cast<const char*>(haystack.data());
+    if (haystack.empty()) { // memchr requires not null
+      return false;
+    }
+    return std::memchr(ptr, needle, haystack.size()) != nullptr;
+  } else if constexpr (sizeof(T) == sizeof(wchar_t)) {
+    auto* ptr = reinterpret_cast<const wchar_t*>(haystack.data());
+    if (haystack.empty()) { // wmemchr requires not null
+      return false;
+    }
+    return std::wmemchr(ptr, needle, haystack.size()) != nullptr;
+  } else {
+    return std::any_of(haystack.begin(), haystack.end(), [needle](T x) {
+      return x == needle;
+    });
+  }
+}
+
+template <typename T>
+constexpr bool hasHandwrittenSimdContains() {
+  return std::is_same_v<T, std::uint8_t> &&
+      !std::is_same_v<SimdCharPlatform, void>;
+}
+
+template <typename T>
+FOLLY_ALWAYS_INLINE bool simdContainsImplHandwritten(
+    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_ALWAYS_INLINE bool simdContainsImpl(
+    folly::span<const T> haystack, T needle) {
+  if constexpr (hasHandwrittenSimdContains<T>()) {
+    return simdContainsImplHandwritten(haystack, needle);
+  } else {
+    return simdContainsImplStd(haystack, needle);
+  }
+}
+
+} // namespace folly::simd_detail
diff --git a/folly/algorithm/simd/test/BUCK b/folly/algorithm/simd/test/BUCK
index 3844ae36967..db05e47f8df 100644
--- a/folly/algorithm/simd/test/BUCK
+++ b/folly/algorithm/simd/test/BUCK
@@ -36,3 +36,14 @@ cpp_benchmark(
         "//folly/init:init",
     ],
 )
+
+cpp_unittest(
+    name = "simd_contains_test",
+    srcs = ["SimdContainsTest.cpp"],
+    headers = [],
+    deps = [
+        "//folly/algorithm/simd:simd_contains",
+        "//folly/algorithm/simd/detail:simd_contains_impl",
+        "//folly/portability:gtest",
+    ],
+)
diff --git a/folly/algorithm/simd/test/SimdContainsTest.cpp b/folly/algorithm/simd/test/SimdContainsTest.cpp
new file mode 100644
index 00000000000..cbae45c33c1
--- /dev/null
+++ b/folly/algorithm/simd/test/SimdContainsTest.cpp
@@ -0,0 +1,78 @@
+/*
+ * 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/SimdContains.h>
+
+#include <folly/algorithm/simd/detail/SimdContainsImpl.h>
+
+#include <folly/portability/GTest.h>
+
+namespace folly {
+
+struct SimdContainsTest : ::testing::Test {};
+
+template <typename T>
+void testSimdContainsVerify(std::span<T> haystack, T needle, bool expected) {
+  bool actual1 = simd_contains(haystack, needle);
+  ASSERT_EQ(expected, actual1);
+
+  auto const_haystack = folly::static_span_cast<const T>(haystack);
+
+  if constexpr (
+      std::is_same_v<T, std::uint8_t> || std::is_same_v<T, std::uint16_t> ||
+      std::is_same_v<T, std::uint32_t> || std::is_same_v<T, std::uint64_t>) {
+    bool actual2 = simd_detail::simdContainsImplStd(const_haystack, needle);
+    ASSERT_EQ(expected, actual2) << " haystack.size(): " << haystack.size();
+  }
+
+  if constexpr (std::is_same_v<T, std::uint8_t>) {
+    bool actual3 =
+        simd_detail::simdContainsImplHandwritten(const_haystack, needle);
+    ASSERT_EQ(expected, actual3) << " haystack.size(): " << haystack.size();
+  }
+}
+
+template <typename T>
+void testSimdContains() {
+  for (std::size_t size = 0; size != 100; ++size) {
+    std::vector<T> buf(size, T{0});
+    for (std::size_t offset = 0; offset != std::min(32UL, buf.size());
+         ++offset) {
+      folly::span searching(buf.begin() + offset, buf.end());
+      T needle{1};
+      testSimdContainsVerify(searching, needle, /*expected*/ false);
+
+      for (auto& x : searching) {
+        x = needle;
+        testSimdContainsVerify(searching, needle, /*expected*/ true);
+        x = 0;
+      }
+    }
+  }
+}
+
+TEST_F(SimdContainsTest, AllTypes) {
+  testSimdContains<std::int8_t>();
+  testSimdContains<std::int16_t>();
+  testSimdContains<std::int32_t>();
+  testSimdContains<std::int64_t>();
+  testSimdContains<std::uint8_t>();
+  testSimdContains<std::uint16_t>();
+  testSimdContains<std::uint32_t>();
+  testSimdContains<std::uint64_t>();
+}
+
+} // namespace folly