From 23637ca05baa79e07bf214535fcfc2b892d6d215 Mon Sep 17 00:00:00 2001
From: Benjamin Kramer <benny.kra@googlemail.com>
Date: Thu, 11 Aug 2022 10:38:12 +0200
Subject: [PATCH] [mlir][sparse] Use the correct ABI on x86 and re-enable tests

c7ec6e19d5446a448f888b33f66316cf2ec6ecae made LLVM adhere to the x86
psABI and pass bf16 in SSE registers instead of GPRs. This breaks the
custom versions of runtime functions we have for bf16 conversion. A
great fix for this would be to use __bf16 types instead which carry the
right ABI, but that type isn't widely available.

Instead just pretend it's a 32 bit float on the ABI boundary and
carefully cast it to the right type.

Fixes #57042

(cherry picked from commit f695554a2a5550ae40da35af9ac22bfcca5db09a)
---
 mlir/lib/ExecutionEngine/Float16bits.cpp      | 47 +++++++++++++------
 .../SparseTensor/CPU/dense_output_bf16.mlir   |  3 --
 .../SparseTensor/CPU/sparse_sum_bf16.mlir     |  3 --
 3 files changed, 32 insertions(+), 21 deletions(-)

diff --git a/mlir/lib/ExecutionEngine/Float16bits.cpp b/mlir/lib/ExecutionEngine/Float16bits.cpp
index c163316913cf..5c3fb61cea2e 100644
--- a/mlir/lib/ExecutionEngine/Float16bits.cpp
+++ b/mlir/lib/ExecutionEngine/Float16bits.cpp
@@ -13,6 +13,7 @@
 
 #include "mlir/ExecutionEngine/Float16bits.h"
 #include <cmath>
+#include <cstring>
 
 namespace {
 
@@ -146,30 +147,46 @@ std::ostream &operator<<(std::ostream &os, const bf16 &d) {
   return os;
 }
 
-// Provide a float->bfloat conversion routine in case the runtime doesn't have
-// one.
-extern "C" uint16_t
+// Mark these symbols as weak so they don't conflict when compiler-rt also
+// defines them.
+#define ATTR_WEAK
 #ifdef __has_attribute
 #if __has_attribute(weak) && !defined(__MINGW32__) && !defined(__CYGWIN__) &&  \
     !defined(_WIN32)
-    __attribute__((__weak__))
+#undef ATTR_WEAK
+#define ATTR_WEAK __attribute__((__weak__))
 #endif
 #endif
-    __truncsfbf2(float f) {
-  return float2bfloat(f);
+
+#if defined(__x86_64__)
+// On x86 bfloat16 is passed in SSE2 registers. Since both float and _Float16
+// are passed in the same register we can use the wider type and careful casting
+// to conform to x86_64 psABI. This only works with the assumption that we're
+// dealing with little-endian values passed in wider registers.
+using BF16ABIType = float;
+#else
+// Default to uint16_t if we have nothing else.
+using BF16ABIType = uint16_t;
+#endif
+
+// Provide a float->bfloat conversion routine in case the runtime doesn't have
+// one.
+extern "C" BF16ABIType ATTR_WEAK __truncsfbf2(float f) {
+  uint16_t bf = float2bfloat(f);
+  // The output can be a float type, bitcast it from uint16_t.
+  BF16ABIType ret = 0;
+  std::memcpy(&ret, &bf, sizeof(bf));
+  return ret;
 }
 
 // Provide a double->bfloat conversion routine in case the runtime doesn't have
 // one.
-extern "C" uint16_t
-#ifdef __has_attribute
-#if __has_attribute(weak) && !defined(__MINGW32__) && !defined(__CYGWIN__) &&  \
-    !defined(_WIN32)
-    __attribute__((__weak__))
-#endif
-#endif
-    __truncdfbf2(double d) {
+extern "C" BF16ABIType ATTR_WEAK __truncdfbf2(double d) {
   // This does a double rounding step, but it's precise enough for our use
   // cases.
-  return __truncsfbf2(static_cast<float>(d));
+  uint16_t bf = __truncsfbf2(static_cast<float>(d));
+  // The output can be a float type, bitcast it from uint16_t.
+  BF16ABIType ret = 0;
+  std::memcpy(&ret, &bf, sizeof(bf));
+  return ret;
 }
diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/dense_output_bf16.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/dense_output_bf16.mlir
index fb392365c00b..f776a3d1b8a0 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/dense_output_bf16.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/dense_output_bf16.mlir
@@ -1,6 +1,3 @@
-// FIXME: see #57042
-// UNSUPPORTED: i386, x86_64
-
 // RUN: mlir-opt %s --sparse-compiler | \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \
diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum_bf16.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum_bf16.mlir
index 6f6a47828d1d..f3307c510c4d 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum_bf16.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum_bf16.mlir
@@ -1,6 +1,3 @@
-// FIXME: see #57042
-// UNSUPPORTED: i386, x86_64
-
 // RUN: mlir-opt %s --sparse-compiler | \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \