From 0f9883f804e4f7847e7215950a1eabfe047543d7 Mon Sep 17 00:00:00 2001
From: Edward Chen <18449977+edgchen1@users.noreply.github.com>
Date: Thu, 20 Jul 2023 18:24:57 -0700
Subject: [PATCH] Fix Mac M1 build (#16763)

- Add ifndef `__APPLE__` to skip lines which cause EXC_BAD_INSTRUCTION error.
- Fix floatToHalf/doubleToHalf conversion issue and add tests.
---
 onnxruntime/core/mlas/lib/platform.cpp  |  5 ++++-
 onnxruntime/core/util/math_cpu.cc       | 21 +++++++++++++++++++--
 onnxruntime/test/framework/math_test.cc | 18 ++++++++++++++++++
 3 files changed, 41 insertions(+), 3 deletions(-)

diff --git a/onnxruntime/core/mlas/lib/platform.cpp b/onnxruntime/core/mlas/lib/platform.cpp
index 6446007610..d7c5f9fc67 100644
--- a/onnxruntime/core/mlas/lib/platform.cpp
+++ b/onnxruntime/core/mlas/lib/platform.cpp
@@ -449,10 +449,13 @@ Return Value:
 
 #if defined(_WIN32)
     HasDotProductInstructions = (IsProcessorFeaturePresent(PF_ARM_V82_DP_INSTRUCTIONS_AVAILABLE) != 0);
-#else
+#elif !defined(__APPLE__)  // The next few lines result in an EXC_BAD_INSTRUCTION runtime error on a M1 Mac so we
+                           // disable it there.
     uint64_t isar0_el1;
     asm("mrs %[reg], ID_AA64ISAR0_EL1\n" : [reg] "=r"(isar0_el1) : :);
     HasDotProductInstructions = ((isar0_el1 >> 44) & 0xfu) == 0x1u;
+#else
+    HasDotProductInstructions = MLAS_CPUIDINFO::GetCPUIDInfo().HasArmNeonDot();
 #endif
 
     if (HasDotProductInstructions) {
diff --git a/onnxruntime/core/util/math_cpu.cc b/onnxruntime/core/util/math_cpu.cc
index f11fe5a701..983321593a 100644
--- a/onnxruntime/core/util/math_cpu.cc
+++ b/onnxruntime/core/util/math_cpu.cc
@@ -20,6 +20,7 @@
 #include "core/framework/float16.h"
 
 #include <algorithm>
+#include <type_traits>
 #include "core/common/narrow.h"
 #include "core/mlas/inc/mlas.h"
 #if defined(__GNUC__)
@@ -855,12 +856,28 @@ void Col2imNd<float, CPUMathUtil, StorageOrder::NCHW>(const float* data_col, con
 SPECIALIZED_COPYVECTOR(float)
 #undef SPECIALIZED_COPYVECTOR
 
+// like C++20's std::bit_cast
+// adapted from the example implementation here: https://en.cppreference.com/w/cpp/numeric/bit_cast
+// TODO replace this with std::bit_cast when we move to C++20
+template <typename Dst, typename Src>
+static std::enable_if_t<
+    sizeof(Src) == sizeof(Dst) &&
+        std::is_trivially_copyable_v<Src> &&
+        std::is_trivially_copyable_v<Dst> &&
+        std::is_trivially_constructible_v<Dst>,
+    Dst>
+BitCast(const Src& src) {
+  Dst dst;
+  std::memcpy(&dst, &src, sizeof(dst));
+  return dst;
+}
+
 uint16_t floatToHalf(float f) {
-  return Eigen::half_impl::float_to_half_rtne(f).x;
+  return BitCast<uint16_t>(Eigen::half_impl::float_to_half_rtne(f).x);
 }
 
 uint16_t doubleToHalf(double f) {
-  return Eigen::half_impl::float_to_half_rtne(static_cast<float>(f)).x;
+  return BitCast<uint16_t>(Eigen::half_impl::float_to_half_rtne(static_cast<float>(f)).x);
 }
 
 float halfToFloat(uint16_t h) {
diff --git a/onnxruntime/test/framework/math_test.cc b/onnxruntime/test/framework/math_test.cc
index 41f8d86699..a133bbd503 100644
--- a/onnxruntime/test/framework/math_test.cc
+++ b/onnxruntime/test/framework/math_test.cc
@@ -202,4 +202,22 @@ TEST(MathTest, GemvTrans) {
   }
 }
 
+TEST(MathTest, HalfFloatConversion) {
+  constexpr float original_values[] = {-4.0f, -2.0f, -1.0f, -0.5f, 0.0f, 0.5f, 1.0f, 2.0f, 4.0f};
+  for (const auto original_value : original_values) {
+    const auto half_value = math::floatToHalf(original_value);
+    const auto round_trip_value = math::halfToFloat(half_value);
+    EXPECT_EQ(round_trip_value, original_value);
+  }
+}
+
+TEST(MathTest, HalfDoubleConversion) {
+  constexpr double original_values[] = {-4.0f, -2.0f, -1.0f, -0.5f, 0.0f, 0.5f, 1.0f, 2.0f, 4.0f};
+  for (const auto original_value : original_values) {
+    const auto half_value = math::doubleToHalf(original_value);
+    const auto round_trip_value = static_cast<double>(math::halfToFloat(half_value));
+    EXPECT_EQ(round_trip_value, original_value);
+  }
+}
+
 }  // namespace onnxruntime