From d275d4c20dbc0164fd814eb28bd2498b6dbb59d3 Mon Sep 17 00:00:00 2001
From: Denis Yaroshevskiy <dyaroshev@meta.com>
Date: Wed, 14 Aug 2024 08:07:11 -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.

Differential Revision: D61205292
---
 CMakeLists.txt                                |  3 +-
 folly/algorithm/simd/BUCK                     |  1 +
 folly/algorithm/simd/FindFixed.h              | 14 +--
 folly/algorithm/simd/detail/BUCK              |  8 ++
 folly/algorithm/simd/detail/Traits.h          | 95 +++++++++++++++++++
 folly/algorithm/simd/detail/test/BUCK         | 10 ++
 .../algorithm/simd/detail/test/TraitsTest.cpp | 88 +++++++++++++++++
 7 files changed, 207 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 c58dedc4ef2..8f34c0f447c 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -646,8 +646,9 @@ if (BUILD_TESTS OR BUILD_BENCHMARKS)
   folly_define_tests(
     DIRECTORY algorithm/simd/detail/test/
       TEST simd_any_of_test SOURCES SimdAnyOfTest.cpp
-      TEST unroll_utils_test SOURCES UnrollUtilsTest.cpp
       TEST simd_for_each_test SOURCES SimdForEachTest.cpp
+      TEST traits_test SOURCES TraitsTest.cpp
+      TEST unroll_utils_test SOURCES UnrollUtilsTest.cpp
 
     DIRECTORY algorithm/simd/test/
       TEST find_fixed_test SOURCES FindFixedTest.cpp
diff --git a/folly/algorithm/simd/BUCK b/folly/algorithm/simd/BUCK
index c1057ab4e84..a660fce85ab 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..310752e9729 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::asSimdFriendly(where), detail::asSimdFriendly(x));
   }
 }
 
diff --git a/folly/algorithm/simd/detail/BUCK b/folly/algorithm/simd/detail/BUCK
index c7501e203a1..eb4cf8a6a2e 100644
--- a/folly/algorithm/simd/detail/BUCK
+++ b/folly/algorithm/simd/detail/BUCK
@@ -36,6 +36,14 @@ cpp_library(
     ],
 )
 
+cpp_library(
+    name = "traits",
+    headers = ["Traits.h"],
+    exported_deps = [
+        "//folly:memory",
+    ],
+)
+
 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..34bdf990c3e
--- /dev/null
+++ b/folly/algorithm/simd/detail/Traits.h
@@ -0,0 +1,95 @@
+/*
+ * 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 <concepts>
+#include <span>
+#include <type_traits>
+
+namespace folly::detail {
+
+struct regular_nonesuch {};
+
+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 {
+      return regular_nonesuch{};
+    }
+  } 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{};
+    } else {
+      return regular_nonesuch{};
+    }
+  } else {
+    return regular_nonesuch{};
+  }
+}
+
+template <bool Cond, typename To>
+using add_const_if_t = std::conditional_t<Cond, const To, To>;
+
+template <typename T>
+concept has_simd_friendly_equivalent =
+    !std::
+        is_same_v<regular_nonesuch, decltype(findSimdFriendlyEquivalent<T>())>;
+
+template <has_simd_friendly_equivalent T>
+using simd_friendly_equivalent_t = add_const_if_t<
+    std::is_const_v<T>,
+    decltype(findSimdFriendlyEquivalent<std::remove_const_t<T>>())>;
+
+template <has_simd_friendly_equivalent T>
+  requires std::integral<T>
+using integral_simd_friendly_equivalent = simd_friendly_equivalent_t<T>;
+
+template <has_simd_friendly_equivalent T, std::size_t Extend>
+auto asSimdFriendly(std::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);
+}
+
+} // 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..a42163f6955
--- /dev/null
+++ b/folly/algorithm/simd/detail/test/TraitsTest.cpp
@@ -0,0 +1,88 @@
+/*
+ * 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 {
+
+TEST(FollySimdTraits, SimdFriendlyEquivalent) {
+  // 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
+  {
+    auto call = []<typename T>(T) -> simd_friendly_equivalent_t<T> {
+      return {};
+    };
+
+    static_assert(std::invocable<decltype(call), int>);
+
+    struct NotSimdFriendly {};
+    static_assert(!std::invocable<decltype(call), NotSimdFriendly>);
+  }
+}
+
+TEST(FollySimdTraits, AsSimdFriendly) {
+  enum SomeEnum : int { Foo = 1, Bar, Baz };
+
+  static_assert(asSimdFriendly(SomeEnum::Foo) == 1);
+
+  std::array arr{SomeEnum::Foo, SomeEnum::Bar, SomeEnum::Baz};
+  std::span<int, 3> castSpan = asSimdFriendly(std::span(arr));
+  ASSERT_THAT(castSpan, testing::ElementsAre(1, 2, 3));
+}
+
+} // namespace folly::detail