From 7975e30ce9b9ef1282fea91626daffa09b2d72ea Mon Sep 17 00:00:00 2001
From: Denis Yaroshevskiy <dyaroshev@meta.com>
Date: Mon, 23 Sep 2024 11:50:53 -0700
Subject: [PATCH] folly simd-friendly (#2279)

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

extracting the common "simd-friendly" type helpers.

Reviewed By: Gownta

Differential Revision: D61205292

fbshipit-source-id: 17551ed2baa168a6b1e5c7dda15945263d97d5fe
---
 CMakeLists.txt                                |   4 +-
 folly/algorithm/simd/BUCK                     |   1 +
 folly/algorithm/simd/FindFixed.h              |  14 +-
 folly/algorithm/simd/detail/BUCK              |  10 ++
 folly/algorithm/simd/detail/Traits.h          | 100 +++++++++++++
 folly/algorithm/simd/detail/test/BUCK         |  10 ++
 .../algorithm/simd/detail/test/TraitsTest.cpp | 139 ++++++++++++++++++
 7 files changed, 266 insertions(+), 12 deletions(-)
 create mode 100644 folly/algorithm/simd/detail/Traits.h
 create mode 100644 folly/algorithm/simd/detail/test/TraitsTest.cpp

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 885990ecc7c..c30fc9938ed 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -648,8 +648,10 @@ if (BUILD_TESTS OR BUILD_BENCHMARKS)
   folly_define_tests(
     DIRECTORY algorithm/simd/detail/test/
       TEST algorithm_simd_detail_simd_any_of_test SOURCES SimdAnyOfTest.cpp
-      TEST algorithm_simd_detail_unroll_utils_test SOURCES UnrollUtilsTest.cpp
       TEST algorithm_simd_detail_simd_for_each_test SOURCES SimdForEachTest.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
diff --git a/folly/algorithm/simd/BUCK b/folly/algorithm/simd/BUCK
index 9e95c669b97..1343e8f9b73 100644
--- a/folly/algorithm/simd/BUCK
+++ b/folly/algorithm/simd/BUCK
@@ -19,5 +19,6 @@ cpp_library(
     exported_deps = [
         ":movemask",
         "//folly:portability",
+        "//folly/algorithm/simd/detail:traits",
     ],
 )
diff --git a/folly/algorithm/simd/FindFixed.h b/folly/algorithm/simd/FindFixed.h
index c76c8022341..547b02a6b3a 100644
--- a/folly/algorithm/simd/FindFixed.h
+++ b/folly/algorithm/simd/FindFixed.h
@@ -28,6 +28,7 @@
 
 #include <folly/Portability.h>
 #include <folly/algorithm/simd/Movemask.h>
+#include <folly/algorithm/simd/detail/Traits.h>
 
 #if FOLLY_X64
 #include <immintrin.h>
@@ -82,11 +83,6 @@ constexpr std::optional<std::size_t> findFixed(std::span<const T, N> where, U x)
 // implementation ---------------------------------------------------------
 
 namespace find_fixed_detail {
-template <typename U, typename T, std::size_t N>
-std::optional<std::size_t> findFixedCast(std::span<const T, N>& where, T x) {
-  std::span<const U, N> whereU{reinterpret_cast<const U*>(where.data()), N};
-  return findFixed(whereU, static_cast<U>(x));
-}
 
 template <typename T>
 constexpr std::optional<std::size_t> findFixedConstexpr(
@@ -295,13 +291,9 @@ constexpr std::optional<std::size_t> findFixed(std::span<const T, N> where, U x)
     return findFixed(where, static_cast<T>(x));
   } else if (std::is_constant_evaluated()) {
     return find_fixed_detail::findFixedConstexpr(std::span<const T>(where), x);
-  } else if constexpr (std::is_enum_v<T>) {
-    return find_fixed_detail::findFixedCast<std::underlying_type_t<T>>(
-        where, x);
-  } else if constexpr (std::is_signed_v<T>) {
-    return find_fixed_detail::findFixedCast<std::make_unsigned_t<T>>(where, x);
   } else {
-    return find_fixed_detail::findFixedDispatch(where, x);
+    return find_fixed_detail::findFixedDispatch(
+        detail::asSimdFriendlyUint(where), detail::asSimdFriendlyUint(x));
   }
 }
 
diff --git a/folly/algorithm/simd/detail/BUCK b/folly/algorithm/simd/detail/BUCK
index c7501e203a1..53a172a2e6b 100644
--- a/folly/algorithm/simd/detail/BUCK
+++ b/folly/algorithm/simd/detail/BUCK
@@ -36,6 +36,16 @@ cpp_library(
     ],
 )
 
+cpp_library(
+    name = "traits",
+    headers = ["Traits.h"],
+    exported_deps = [
+        "//folly:memory",
+        "//folly:traits",
+        "//folly/container:span",
+    ],
+)
+
 cpp_library(
     name = "unroll_utils",
     headers = ["UnrollUtils.h"],
diff --git a/folly/algorithm/simd/detail/Traits.h b/folly/algorithm/simd/detail/Traits.h
new file mode 100644
index 00000000000..ee2ef3d5105
--- /dev/null
+++ b/folly/algorithm/simd/detail/Traits.h
@@ -0,0 +1,100 @@
+/*
+ * 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/Memory.h>
+#include <folly/Traits.h>
+#include <folly/container/span.h>
+
+#include <concepts>
+#include <type_traits>
+
+namespace folly::detail {
+
+template <typename T>
+auto findSimdFriendlyEquivalent() {
+  if constexpr (std::is_enum_v<T>) {
+    return findSimdFriendlyEquivalent<std::underlying_type_t<T>>();
+  } else if constexpr (std::is_floating_point_v<T>) {
+    if constexpr (sizeof(T) == 4) {
+      return float{};
+    } else {
+      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{};
+    }
+  } 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{};
+    }
+  }
+}
+
+template <typename T>
+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< //
+    T,
+    decltype(findSimdFriendlyEquivalent<std::remove_const_t<T>>())>;
+
+template <typename T>
+concept has_integral_simd_friendly_equivalent =
+    has_simd_friendly_equivalent<T> && // have to explicitly specify this for
+                                       // subsumption to work
+    std::integral<simd_friendly_equivalent_t<T>>;
+
+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);
+}
+
+template <has_simd_friendly_equivalent T>
+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);
+}
+
+template <has_simd_friendly_equivalent T>
+constexpr auto asSimdFriendlyUint(T x) {
+  return static_cast<uint_bits_t<sizeof(T) * 8>>(x);
+}
+
+} // namespace folly::detail
diff --git a/folly/algorithm/simd/detail/test/BUCK b/folly/algorithm/simd/detail/test/BUCK
index fe5871aa577..6661c2ce012 100644
--- a/folly/algorithm/simd/detail/test/BUCK
+++ b/folly/algorithm/simd/detail/test/BUCK
@@ -24,6 +24,16 @@ cpp_unittest(
     ],
 )
 
+cpp_unittest(
+    name = "traits_test",
+    srcs = ["TraitsTest.cpp"],
+    deps = [
+        "//folly/algorithm/simd/detail:traits",
+        "//folly/portability:gmock",
+        "//folly/portability:gtest",
+    ],
+)
+
 cpp_unittest(
     name = "unroll_utils_test",
     srcs = [
diff --git a/folly/algorithm/simd/detail/test/TraitsTest.cpp b/folly/algorithm/simd/detail/test/TraitsTest.cpp
new file mode 100644
index 00000000000..5d01d9630df
--- /dev/null
+++ b/folly/algorithm/simd/detail/test/TraitsTest.cpp
@@ -0,0 +1,139 @@
+/*
+ * 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/detail/Traits.h>
+
+#include <folly/portability/GMock.h>
+#include <folly/portability/GTest.h>
+
+namespace folly::detail {
+
+struct FollySimdTraitsTest : testing::Test {};
+
+namespace simd_friendly_equivalent_test {
+
+// ints
+static_assert(
+    std::is_same_v<std::int8_t, simd_friendly_equivalent_t<signed char>>);
+static_assert(
+    std::is_same_v<std::uint8_t, simd_friendly_equivalent_t<unsigned char>>);
+
+static_assert(std::is_same_v<std::int16_t, simd_friendly_equivalent_t<short>>);
+static_assert(
+    std::is_same_v<std::uint16_t, simd_friendly_equivalent_t<unsigned short>>);
+
+static_assert(std::is_same_v<std::int32_t, simd_friendly_equivalent_t<int>>);
+static_assert(
+    std::is_same_v<std::uint32_t, simd_friendly_equivalent_t<unsigned int>>);
+
+static_assert(
+    std::is_same_v<std::int64_t, simd_friendly_equivalent_t<std::int64_t>>);
+static_assert(
+    std::is_same_v<std::uint64_t, simd_friendly_equivalent_t<std::uint64_t>>);
+
+// floats
+static_assert(std::is_same_v<float, simd_friendly_equivalent_t<float>>);
+static_assert(std::is_same_v<double, simd_friendly_equivalent_t<double>>);
+
+// enum
+enum SomeInt {};
+enum class SomeIntClass : std::int32_t {};
+
+static_assert(
+    std::is_same_v<std::uint32_t, simd_friendly_equivalent_t<SomeInt>>);
+static_assert(
+    std::is_same_v<std::int32_t, simd_friendly_equivalent_t<SomeIntClass>>);
+
+// const
+
+static_assert(
+    std::is_same_v<const std::int32_t, simd_friendly_equivalent_t<const int>>);
+
+// sfinae
+constexpr auto sfinae_call =
+    []<typename T>(T) -> simd_friendly_equivalent_t<T> { return {}; };
+
+static_assert(std::invocable<decltype(sfinae_call), int>);
+
+struct NotSimdFriendly {};
+static_assert(!std::invocable<decltype(sfinae_call), NotSimdFriendly>);
+
+} // namespace simd_friendly_equivalent_test
+
+namespace integral_simd_friendly_equivalent_test {
+
+static_assert(std::is_same_v< //
+              std::int8_t,
+              integral_simd_friendly_equivalent<signed char>>);
+
+struct Overloading {
+  constexpr int operator()(auto) { return 0; }
+  constexpr int operator()(has_simd_friendly_equivalent auto) { return 1; }
+  constexpr int operator()(has_integral_simd_friendly_equivalent auto) {
+    return 2;
+  }
+};
+
+// Subsumption tests
+struct NotSimdFriendly {};
+enum class SomeInt {};
+
+static_assert(Overloading{}(NotSimdFriendly{}) == 0);
+static_assert(Overloading{}(float{}) == 1);
+static_assert(Overloading{}(int{}) == 2);
+static_assert(Overloading{}(SomeInt{}) == 2);
+
+} // namespace integral_simd_friendly_equivalent_test
+
+TEST_F(FollySimdTraitsTest, AsSimdFriendly) {
+  enum SomeEnum : int { Foo = 1, Bar, Baz };
+
+  static_assert(asSimdFriendly(SomeEnum::Foo) == 1);
+
+  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));
+}
+
+template <typename T, typename U>
+void isSameTest(const T&, const U&) = delete;
+
+template <typename T>
+void isSameTest(const T&, const T&) {}
+
+template <typename From, typename To>
+void asSimdFriendlyUintTypeTest() {
+  isSameTest(asSimdFriendlyUint(From{}), To{});
+  isSameTest(asSimdFriendlyUint(std::span<From>{}), std::span<To>{});
+  isSameTest(
+      asSimdFriendlyUint(std::span<const From>{}), std::span<const To>{});
+}
+
+TEST_F(FollySimdTraitsTest, AsSimdFriendlyUint) {
+  enum SomeEnum : int { Foo = 1, Bar, Baz };
+
+  static_assert(asSimdFriendlyUint(SomeEnum::Foo) == 1U);
+
+  asSimdFriendlyUintTypeTest<char, std::uint8_t>();
+  asSimdFriendlyUintTypeTest<short, std::uint16_t>();
+  asSimdFriendlyUintTypeTest<int, std::uint32_t>();
+  asSimdFriendlyUintTypeTest<unsigned, std::uint32_t>();
+  asSimdFriendlyUintTypeTest<float, std::uint32_t>();
+  asSimdFriendlyUintTypeTest<int64_t, std::uint64_t>();
+  asSimdFriendlyUintTypeTest<double, std::uint64_t>();
+}
+
+} // namespace folly::detail