support pointers in simd contains

Summary: pointers should be supported. Reviewed By: yfeldblum Differential Revision: D64038733 fbshipit-source-id: 3b3e91831c242b4e4f8b4be4f63b9c4983b2957a
facebook · Oct 10, 2024 · 475d3e4 · 475d3e4
1 parent df3fdab
commit 475d3e4
Show file tree

Hide file tree

Showing 4 changed files with 46 additions and 3 deletions.
diff --git a/folly/algorithm/simd/Contains.h b/folly/algorithm/simd/Contains.h
@@ -94,7 +94,7 @@ struct contains_fn {
     auto castR = detail::asSimdFriendlyUint(folly::span(r));
     using value_type = detail::std_range_value_t<decltype(castR)>;
 
-    auto castX = static_cast<value_type>(x);
+    auto castX = static_cast<value_type>(detail::asSimdFriendlyUint(x));
 
     if constexpr (std::is_same_v<value_type, std::uint8_t>) {
       return detail::containsU8(castR, castX);

diff --git a/folly/algorithm/simd/detail/Traits.h b/folly/algorithm/simd/detail/Traits.h
@@ -30,6 +30,10 @@ template <typename T>
 auto findSimdFriendlyEquivalent() {
   if constexpr (std::is_enum_v<T>) {
     return findSimdFriendlyEquivalent<std::underlying_type_t<T>>();
+  } else if constexpr (std::is_pointer_v<T>) {
+    // We use signed numbers for pointers because x86 support for signed
+    // numbers is better and we can get away with it, in terms of correctness.
+    return int_bits_t<sizeof(T) * 8>{};
   } else if constexpr (std::is_floating_point_v<T>) {
     if constexpr (sizeof(T) == 4) {
       return float{};
@@ -81,7 +85,12 @@ struct AsSimdFriendlyFn {
   template <typename T>
   FOLLY_ERASE constexpr auto operator()(T x) const
       -> simd_friendly_equivalent_scalar_t<T> {
-    return static_cast<simd_friendly_equivalent_scalar_t<T>>(x);
+    using res_t = simd_friendly_equivalent_scalar_t<T>;
+    if constexpr (!std::is_pointer_v<T>) {
+      return static_cast<res_t>(x);
+    } else {
+      return reinterpret_cast<res_t>(x);
+    }
   }
 };
 inline constexpr AsSimdFriendlyFn asSimdFriendly;
@@ -103,7 +112,12 @@ struct AsSimdFriendlyUintFn {
   template <typename T>
   FOLLY_ERASE constexpr auto operator()(T x) const
       -> unsigned_simd_friendly_equivalent_scalar_t<T> {
-    return static_cast<unsigned_simd_friendly_equivalent_scalar_t<T>>(x);
+    using res_t = unsigned_simd_friendly_equivalent_scalar_t<T>;
+    if constexpr (!std::is_pointer_v<T>) {
+      return static_cast<res_t>(x);
+    } else {
+      return reinterpret_cast<res_t>(x);
+    }
   }
 };
 inline constexpr AsSimdFriendlyUintFn asSimdFriendlyUint;

diff --git a/folly/algorithm/simd/detail/test/TraitsTest.cpp b/folly/algorithm/simd/detail/test/TraitsTest.cpp
@@ -76,6 +76,12 @@ static_assert(std::is_same_v<
               const std::int32_t,
               simd_friendly_equivalent_scalar_t<const int>>);
 
+// pointers
+
+static_assert(
+    sizeof(std::int64_t) != sizeof(void*) ||
+    std::is_same_v<std::int64_t, simd_friendly_equivalent_scalar_t<void*>>);
+
 // sfinae
 
 struct sfinae_call {
@@ -169,6 +175,13 @@ TEST_F(SimdTraitsTest, AsSimdFriendly) {
   std::array arr{SomeEnum::Foo, SomeEnum::Bar, SomeEnum::Baz};
   folly::span<int, 3> castSpan = asSimdFriendly(folly::span(arr));
   ASSERT_THAT(castSpan, testing::ElementsAre(1, 2, 3));
+
+  // pointer
+  {
+    auto expected = folly::bit_cast<std::intptr_t>(arr.data());
+    auto actual = asSimdFriendly(arr.data());
+    ASSERT_EQ(expected, actual);
+  }
 }
 
 TEST_F(SimdTraitsTest, AsSimdFriendlyUint) {
@@ -179,6 +192,13 @@ TEST_F(SimdTraitsTest, AsSimdFriendlyUint) {
   std::array arr{SomeEnum::Foo, SomeEnum::Bar, SomeEnum::Baz};
   folly::span<std::uint32_t, 3> castSpan = asSimdFriendlyUint(folly::span(arr));
   ASSERT_THAT(castSpan, testing::ElementsAre(1, 2, 3));
+
+  // pointer
+  {
+    auto expected = folly::bit_cast<std::uintptr_t>(arr.data());
+    auto actual = asSimdFriendlyUint(arr.data());
+    ASSERT_EQ(expected, actual);
+  }
 }
 
 } // namespace folly::simd::detail
diff --git a/folly/algorithm/simd/test/ContainsTest.cpp b/folly/algorithm/simd/test/ContainsTest.cpp
@@ -94,6 +94,8 @@ static_assert( //
 template <typename T>
 struct ContainsTest : ::testing::Test {};
 
+struct ContainsTestSpeicalCases : ::testing::Test {};
+
 using TypesToTest = ::testing::Types<
     std::int8_t,
     std::int16_t,
@@ -149,4 +151,11 @@ TYPED_TEST(ContainsTest, Basic) {
   }
 }
 
+TEST_F(ContainsTestSpeicalCases, Pointers) {
+  std::array ints = {0, 1, 2, 3};
+  std::array ptrs = {&ints[0], &ints[1], &ints[3]};
+  EXPECT_TRUE(folly::simd::contains(ptrs, &ints[1]));
+  EXPECT_FALSE(folly::simd::contains(ptrs, &ints[2]));
+}
+
 } // namespace folly::simd