Pass tests/gpu/jit.cpp with BF16

src/targets/gpu/kernels/include/migraphx/kernels/hip.hpp

@@ -27,6 +27,7 @@
 #ifndef MIGRAPHX_USE_HIPRTC
 #include <hip/hip_runtime.h>
 #include <hip/hip_fp16.h>
+#include <hip/hip_bf16.h>
 #include <hip/math_functions.h>
 #endif
 

src/targets/gpu/kernels/include/migraphx/kernels/math.hpp

@@ -40,6 +40,8 @@
 constexpr float as_float(migraphx::fp8::fp8e4m3fn x) { return x; }
 constexpr float as_float(migraphx::fp8::fp8e5m2 x) { return x; }
 
+constexpr float as_float(migraphx::bf16 x) { return x; }
+
 template <class T>
 constexpr T as_float(T x)
 {
@@ -78,6 +80,12 @@
     auto __device__ name(migraphx::half x, Ts... xs)                   \
         MIGRAPHX_RETURNS(fname(math::as_float(x), math::as_float(xs)...))
 
+// NOLINTNEXTLINE
+#define MIGRAPHX_DEVICE_MATH_BF16(name, fname)                         \
+    template <class... Ts, MIGRAPHX_REQUIRES(not is_any_vec<Ts...>())> \
+    auto __device__ name(migraphx::bf16 x, Ts... xs)                   \
+        MIGRAPHX_RETURNS(fname(math::as_float(x), math::as_float(xs)...))
+
 // NOLINTNEXTLINE
 #define MIGRAPHX_DEVICE_MATH_FP8(name, fname)                                          \
     template <class... Ts, MIGRAPHX_REQUIRES(not is_any_vec<Ts...>())>                 \
@@ -166,6 +174,20 @@
 MIGRAPHX_DEVICE_MATH_FOR(migraphx::half, sin, ::hsin)
 MIGRAPHX_DEVICE_MATH_FOR(migraphx::half, sqrt, ::hsqrt)
 
+// Builtin half functions
+// MIGRAPHX_DEVICE_MATH_FOR(migraphx::bf16, abs, ::__habs)
+// MIGRAPHX_DEVICE_MATH_FOR(migraphx::bf16, ceil, ::hceil)
+// MIGRAPHX_DEVICE_MATH_FOR(migraphx::bf16, cos, ::hcos)
+// MIGRAPHX_DEVICE_MATH_FOR(migraphx::bf16, exp, ::hexp)
+// MIGRAPHX_DEVICE_MATH_FOR(migraphx::bf16, floor, ::hfloor)
+// MIGRAPHX_DEVICE_MATH_FOR(migraphx::bf16, isinf, ::__hisinf)
+// MIGRAPHX_DEVICE_MATH_FOR(migraphx::bf16, isnan, ::__hisnan)
+// MIGRAPHX_DEVICE_MATH_FOR(migraphx::bf16, log, ::hlog)
+// MIGRAPHX_DEVICE_MATH_FOR(migraphx::bf16, log2, ::hlog2)
+// MIGRAPHX_DEVICE_MATH_FOR(migraphx::bf16, rsqrt, ::hrsqrt)
+// MIGRAPHX_DEVICE_MATH_FOR(migraphx::bf16, sin, ::hsin)
+// MIGRAPHX_DEVICE_MATH_FOR(migraphx::bf16, sqrt, ::hsqrt)
+
 // Use float to compute half overload
 MIGRAPHX_DEVICE_MATH_HALF(acos, ::acos)
 MIGRAPHX_DEVICE_MATH_HALF(acosh, ::acosh)
@@ -184,6 +206,34 @@
 MIGRAPHX_DEVICE_MATH_HALF(tanh, ::tanh)
 MIGRAPHX_DEVICE_MATH_HALF(fmod, ::fmod)
 
+// Use float to compute bf16 overload
+MIGRAPHX_DEVICE_MATH_BF16(abs, ::abs)
+MIGRAPHX_DEVICE_MATH_BF16(acos, ::acos)
+MIGRAPHX_DEVICE_MATH_BF16(acosh, ::acosh)
+MIGRAPHX_DEVICE_MATH_BF16(asin, ::asin)
+MIGRAPHX_DEVICE_MATH_BF16(asinh, ::asinh)
+MIGRAPHX_DEVICE_MATH_BF16(atan, ::atan)
+MIGRAPHX_DEVICE_MATH_BF16(atanh, ::atanh)
+MIGRAPHX_DEVICE_MATH_BF16(ceil, ::ceil)
+MIGRAPHX_DEVICE_MATH_BF16(cos, ::cos)
+MIGRAPHX_DEVICE_MATH_BF16(cosh, ::cosh)
+MIGRAPHX_DEVICE_MATH_BF16(erf, ::erf)
+MIGRAPHX_DEVICE_MATH_BF16(exp, ::exp)
+MIGRAPHX_DEVICE_MATH_BF16(floor, ::floor)
+MIGRAPHX_DEVICE_MATH_BF16(isnan, ::isnan)
+MIGRAPHX_DEVICE_MATH_BF16(log, ::log)
+MIGRAPHX_DEVICE_MATH_BF16(log2, ::log2)
+MIGRAPHX_DEVICE_MATH_BF16(pow, ::pow)
+MIGRAPHX_DEVICE_MATH_BF16(remainder, ::remainder)
+MIGRAPHX_DEVICE_MATH_BF16(round, ::round)
+MIGRAPHX_DEVICE_MATH_BF16(rsqrt, ::rsqrt)
+MIGRAPHX_DEVICE_MATH_BF16(sin, ::sin)
+MIGRAPHX_DEVICE_MATH_BF16(sinh, ::sinh)
+MIGRAPHX_DEVICE_MATH_BF16(sqrt, ::sqrt)
+MIGRAPHX_DEVICE_MATH_BF16(tan, ::tan)
+MIGRAPHX_DEVICE_MATH_BF16(tanh, ::tanh)
+MIGRAPHX_DEVICE_MATH_BF16(fmod, ::fmod)
+
 // use float to compute fp8 overload
 MIGRAPHX_DEVICE_MATH_FP8(abs, ::abs)
 MIGRAPHX_DEVICE_MATH_FP8(acos, ::acos)
@@ -242,8 +292,10 @@
 MIGRAPHX_DEVICE_MATH_BINARY_FOR(float, min, ::min)
 MIGRAPHX_DEVICE_MATH_BINARY_FOR(double, max, ::max)
 MIGRAPHX_DEVICE_MATH_BINARY_FOR(double, min, ::min)
 MIGRAPHX_DEVICE_MATH_BINARY_FOR(migraphx::half, max, ::__hmax)
 MIGRAPHX_DEVICE_MATH_BINARY_FOR(migraphx::half, min, ::__hmin)
+// MIGRAPHX_DEVICE_MATH_BINARY_FOR(migraphx::bf16, max, ::__hmax)
+// MIGRAPHX_DEVICE_MATH_BINARY_FOR(migraphx::bf16, min, ::__hmin)
 
 template <class T, MIGRAPHX_REQUIRES(not is_any_vec<T>())>
 constexpr auto max(const T& a, const T& b)

src/targets/gpu/kernels/include/migraphx/kernels/type_traits.hpp

@@ -246,6 +246,8 @@ constexpr T numeric_max()
         return __FLT_MAX__;
     else if constexpr(is_same<T, migraphx::half>{})
         return __FLT16_MAX__;
+    else if constexpr(is_same<T, migraphx::bf16>{})
+        return 338953138925153547590470800371487866880.000000;
     else
         return 0;
 }

src/targets/gpu/kernels/include/migraphx/kernels/types.hpp

@@ -76,7 +76,7 @@ using vec = T __attribute__((ext_vector_type(N)));
 
 using half  = _Float16;
 using half2 = migraphx::vec<half, 2>;
-
+using bf16 = __bf16;
 } // namespace migraphx
 
 #endif

test/gpu/jit.cpp

@@ -408,13 +408,12 @@ TEST_CASE(assert_type_min_max)
         if(contains({migraphx::shape::bool_type,
                      migraphx::shape::fp8e4m3fnuz_type,
                      migraphx::shape::tuple_type,
-                     migraphx::shape::bf16_type},
+                     },
                     t))
             continue;
         auto name = migraphx::shape::cpp_type(t);
-        if(t == migraphx::shape::half_type)
+        if((t == migraphx::shape::half_type) or (t == migraphx::shape::bf16_type))
             name.insert(0, "migraphx::");
-
         migraphx::shape::visit(t, [&](auto as) {
             std::string min = "";
             std::string max = "";