Making fma_relu accept the bfloat16 class

Author: Hugh Delaney

sycl/include/sycl/ext/oneapi/experimental/builtins.hpp

@@ -16,6 +16,8 @@
 
 #include <CL/__spirv/spirv_ops.hpp>
 
+#include "bfloat16.hpp"
+
 // TODO Decide whether to mark functions with this attribute.
 #define __NOEXC /*noexcept*/
 
@@ -37,11 +39,18 @@ namespace experimental {
 // fma_relu returns a * b + c > 0 ? a * b + c : 0
 template <typename T>
 sycl::detail::enable_if_t<sycl::detail::is_genfloath<T>::value ||
-                              sycl::detail::is_ugenshort<T>::value ||
-                              sycl::detail::is_ugenint<T>::value,
+                              sycl::detail::is_ugenint<T>::value ||
+                              std::is_same<T, bfloat16>::value,
                           T>
 fma_relu(T a, T b, T c) __NOEXC {
-  return __sycl_std::__invoke_fma_relu<T>(a, b, c);
+  if constexpr (std::is_same<T, bfloat16>::value) {
+    uint16_t tmp = __sycl_std::__invoke_fma_relu<uint16_t>(
+        reinterpret_cast<uint16_t &>(a), reinterpret_cast<uint16_t &>(b),
+        reinterpret_cast<uint16_t &>(c));
+    return reinterpret_cast<bfloat16 &>(tmp);
+  } else {
+    return __sycl_std::__invoke_fma_relu<T>(a, b, c);
+  }
 }
 
 // Provides functionality to print data from kernels in a C way:
@@ -53,9 +62,9 @@ fma_relu(T a, T b, T c) __NOEXC {
 // Please refer to corresponding section in OpenCL C specification to find
 // information about format string and its differences from standard C rules.
 //
-// This function is placed under 'experimental' namespace on purpose, because it
-// has too much caveats you need to be aware of before using it. Please find
-// them below and read carefully before using it:
+// This function is placed under 'experimental' namespace on purpose, because
+// it has too much caveats you need to be aware of before using it. Please
+// find them below and read carefully before using it:
 //
 // - According to the OpenCL spec, the format string must be
 // resolvable at compile time i.e. cannot be dynamically created by the
@@ -65,19 +74,19 @@ fma_relu(T a, T b, T c) __NOEXC {
 // address space. The constant address space declarations might get "tricky",
 // see test/built-ins/printf.cpp for examples.
 // In simple cases (compile-time known string contents, direct declaration of
-// the format literal inside the printf call, etc.), the compiler should handle
-// the automatic address space conversion.
+// the format literal inside the printf call, etc.), the compiler should
+// handle the automatic address space conversion.
 // FIXME: Once the extension to generic address space is fully supported, the
 // constant AS version may need to be deprecated.
 //
-// - The format string is interpreted according to the OpenCL C spec, where all
-// data types has fixed size, opposed to C++ types which doesn't guarantee
+// - The format string is interpreted according to the OpenCL C spec, where
+// all data types has fixed size, opposed to C++ types which doesn't guarantee
 // the exact width of particular data types (except, may be, char). This might
 // lead to unexpected result, for example: %ld in OpenCL C means that printed
-// argument has 'long' type which is 64-bit wide by the OpenCL C spec. However,
-// by C++ spec long is just at least 32-bit wide, so, you need to ensure (by
-// performing a cast, for example) that if you use %ld specifier, you pass
-// 64-bit argument to the cl::sycl::experimental::printf
+// argument has 'long' type which is 64-bit wide by the OpenCL C spec.
+// However, by C++ spec long is just at least 32-bit wide, so, you need to
+// ensure (by performing a cast, for example) that if you use %ld specifier,
+// you pass 64-bit argument to the cl::sycl::experimental::printf
 //
 // - OpenCL spec defines several additional features, like, for example, 'v'
 // modifier which allows to print OpenCL vectors: note that these features are