[llvm-branch-commits] [mlir] 23637ca - [mlir][sparse] Use the correct ABI on x86 and re-enable tests

[Tobias Hieta via llvm-branch-commits]

Author: Benjamin Kramer
Date: 2022-08-12T08:39:37+02:00
New Revision: 23637ca05baa79e07bf214535fcfc2b892d6d215

URL: https://github.com/llvm/llvm-project/commit/23637ca05baa79e07bf214535fcfc2b892d6d215
DIFF: https://github.com/llvm/llvm-project/commit/23637ca05baa79e07bf214535fcfc2b892d6d215.diff

LOG: [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)

Added: 
    

Modified: 
    mlir/lib/ExecutionEngine/Float16bits.cpp
    mlir/test/Integration/Dialect/SparseTensor/CPU/dense_output_bf16.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum_bf16.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/ExecutionEngine/Float16bits.cpp b/mlir/lib/ExecutionEngine/Float16bits.cpp
index c163316913cf5..5c3fb61cea2e2 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
+
+#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
-    __truncsfbf2(float f) {
-  return float2bfloat(f);
+
+// 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 fb392365c00bc..f776a3d1b8a09 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 6f6a47828d1d9..f3307c510c4d4 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  \