Enable fp8e5m2fnuz type

src/api/include/migraphx/migraphx.h

@@ -48,7 +48,8 @@
     m(fp8e4m3fnuz_type, migraphx::fp8::fp8e4m3fnuz) \
     m(fp8e4m3fn_type, migraphx::fp8::fp8e4m3fn) \
     m(fp8e5m2_type, migraphx::fp8::fp8e5m2) \
-    m(bf16_type, bf16)
+    m(bf16_type, bf16) \
+    m(fp8e5m2fnuz_type, migraphx::fp8::fp8e5m2fnuz)
 // clang-format on
 
 #ifdef __cplusplus

src/include/migraphx/fp8_types.hpp

@@ -28,8 +28,10 @@ namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
 struct fp8_types
 {
-    const std::set<shape::type_t> types = {
-        shape::fp8e4m3fnuz_type, shape::fp8e4m3fn_type, shape::fp8e5m2_type};
+    const std::set<shape::type_t> types = {shape::fp8e4m3fnuz_type,
+                                           shape::fp8e5m2fnuz_type,
+                                           shape::fp8e4m3fn_type,
+                                           shape::fp8e5m2_type};
 
     std::set<shape::type_t> get() const { return types; }
 };

src/include/migraphx/generic_float.hpp

@@ -161,8 +161,8 @@ struct __attribute__((packed, may_alias)) generic_float
             }
             else
             {
-                type shift         = 0;
-                f.mantissa         = mantissa;
+                type shift = 0;
+                f.mantissa = mantissa;
 
                 if(MantissaSize < float32_parts::mantissa_width())
                 {

src/include/migraphx/op/quant_convolution.hpp

@@ -31,6 +31,7 @@
 #include <migraphx/config.hpp>
 #include <migraphx/convolution.hpp>
 #include <migraphx/value.hpp>
+#include <migraphx/fp8_types.hpp>
 #include <cmath>
 #include <utility>
 
@@ -89,8 +90,8 @@ struct quant_convolution
 
         // all input type must be int8_type or fp8 types
         // output should be float_type
-        std::set<migraphx::shape::type_t> supported_types = {
-            shape::int8_type, shape::fp8e4m3fnuz_type, shape::fp8e4m3fn_type, shape::fp8e5m2_type};
+        std::set<migraphx::shape::type_t> supported_types = fp8_types{}.get();
+        supported_types.insert(shape::int8_type);
         if(not contains(supported_types, t))
         {
             MIGRAPHX_THROW("QUANT_CONVOLUTION: only accept input and weights of type int8 or fp8");

src/include/migraphx/op/quant_dot.hpp

@@ -45,11 +45,9 @@ struct quant_dot
         const shape& a = inputs.at(0);
         const shape& b = inputs.at(1);
         auto t         = a.type();
-        std::set<migraphx::shape::type_t> supported_types = {shape::int8_type,
-                                                             shape::uint8_type,
-                                                             shape::fp8e4m3fnuz_type,
-                                                             shape::fp8e4m3fn_type,
-                                                             shape::fp8e5m2_type};
+        std::set<migraphx::shape::type_t> supported_types = fp8_types{}.get();
+        supported_types.insert(shape::int8_type);
+        supported_types.insert(shape::uint8_type);
         if(not contains(supported_types, t))
         {
             MIGRAPHX_THROW("QUANT_DOT: only support data type int8_t, uint8_t and fp8 types");

src/include/migraphx/shape.hpp

@@ -66,8 +66,8 @@ struct MIGRAPHX_EXPORT shape
     m(fp8e4m3fnuz_type, migraphx::fp8::fp8e4m3fnuz) \
     m(fp8e4m3fn_type, migraphx::fp8::fp8e4m3fn) \
     m(fp8e5m2_type, migraphx::fp8::fp8e5m2) \
-    m(bf16_type, bf16)
-    // clang-format on
+    m(bf16_type, bf16) \
+    m(fp8e5m2fnuz_type, migraphx::fp8::fp8e5m2fnuz) // clang-format on
 
 #define MIGRAPHX_SHAPE_GENERATE_ENUM_TYPES(x, t) x,
     enum type_t

src/include/migraphx/type_traits.hpp

@@ -62,6 +62,10 @@ MIGRAPHX_DETAIL_EXTEND_TRAIT_FOR(is_floating_point, migraphx::fp8::fp8e4m3fnuz)
 MIGRAPHX_DETAIL_EXTEND_TRAIT_FOR(is_signed, migraphx::fp8::fp8e4m3fnuz)
 MIGRAPHX_DETAIL_EXTEND_TRAIT_FOR(is_arithmetic, migraphx::fp8::fp8e4m3fnuz)
 
+MIGRAPHX_DETAIL_EXTEND_TRAIT_FOR(is_floating_point, migraphx::fp8::fp8e5m2fnuz)
+MIGRAPHX_DETAIL_EXTEND_TRAIT_FOR(is_signed, migraphx::fp8::fp8e5m2fnuz)
+MIGRAPHX_DETAIL_EXTEND_TRAIT_FOR(is_arithmetic, migraphx::fp8::fp8e5m2fnuz)
+
 MIGRAPHX_DETAIL_EXTEND_TRAIT_FOR(is_floating_point, migraphx::fp8::fp8e4m3fn)
 MIGRAPHX_DETAIL_EXTEND_TRAIT_FOR(is_signed, migraphx::fp8::fp8e4m3fn)
 MIGRAPHX_DETAIL_EXTEND_TRAIT_FOR(is_arithmetic, migraphx::fp8::fp8e4m3fn)

src/py/migraphx_py.cpp

@@ -158,6 +158,17 @@ struct npy_format_descriptor<migraphx::fp8::fp8e4m3fnuz>
     static constexpr auto name() { return _("fp8e4m3fnuz"); }
 };
 
+template <>
+struct npy_format_descriptor<migraphx::fp8::fp8e5m2fnuz>
+{
+    static std::string format()
+    {
+        // TODO: no standard format in numpy for fp8
+        return "z";
+    }
+    static constexpr auto name() { return _("fp8e5m2fnuz"); }
+};
+
 template <>
 struct npy_format_descriptor<migraphx::fp8::fp8e4m3fn>
 {

src/simplify_qdq.cpp

@@ -35,6 +35,7 @@
 #include <migraphx/op/dot.hpp>
 #include <migraphx/op/quant_dot.hpp>
 #include <migraphx/register_op.hpp>
+#include <migraphx/fp8_types.hpp>
 
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
@@ -143,10 +144,8 @@ struct match_find_quantizable_ops
         auto zp1    = r.instructions["zp1"];
         auto zp2    = r.instructions["zp2"];
         // Only INT8 or FP8 type currently supported
-        std::set<migraphx::shape::type_t> supported_types = {migraphx::shape::fp8e4m3fnuz_type,
-                                                             migraphx::shape::fp8e4m3fn_type,
-                                                             migraphx::shape::fp8e5m2_type,
-                                                             migraphx::shape::int8_type};
+        std::set<migraphx::shape::type_t> supported_types = fp8_types{}.get();
+        supported_types.insert(migraphx::shape::int8_type);
         if(not contains(supported_types, dq1->inputs().front()->get_shape().type()) or
            not contains(supported_types, dq2->inputs().front()->get_shape().type()))
             return;

src/targets/gpu/fuse_mlir.cpp

@@ -301,8 +301,8 @@ auto is_mlir_conv(mlir_mode mode)
         // Avoid MLIR assertion: Index < Length && "Invalid index!"
         if(ins->get_shape().lens().size() != 4 and group > 1)
             return false;
-        std::set<shape::type_t> supported_types = {
-            shape::fp8e4m3fnuz_type, shape::fp8e4m3fn_type, shape::fp8e5m2_type, shape::int8_type};
+        std::set<shape::type_t> supported_types = fp8_types{}.get();
+        supported_types.insert(shape::int8_type);
         if(contains(supported_types, input.type()))
             return true;
         if(mode == mlir_mode::all)
@@ -366,6 +366,7 @@ bool is_pointwise_op_supported_by_mlir(const instruction& i)
     const std::initializer_list<type_t> allowed_types = {type_t::float_type,
                                                          type_t::half_type,
                                                          type_t::fp8e4m3fnuz_type,
+                                                         type_t::fp8e5m2fnuz_type,
                                                          type_t::fp8e4m3fn_type,
                                                          type_t::fp8e5m2_type,
                                                          type_t::int8_type,
@@ -413,6 +414,7 @@ bool is_pointwise_op_supported_by_mlir(const instruction& i)
     std::set<shape::type_t> float_types = {type_t::float_type,
                                            type_t::half_type,
                                            type_t::fp8e4m3fnuz_type,
+                                           type_t::fp8e5m2fnuz_type,
                                            type_t::fp8e4m3fn_type,
                                            type_t::fp8e5m2_type};
     bool is_float                       = contains(float_types, result_type);
@@ -442,8 +444,12 @@ bool is_reduce_op_supported_by_mlir(const instruction& i)
     using type_t                                      = shape::type_t;
     const auto& name                                  = i.name();
     const auto result_type                            = i.get_shape().type();
-    const std::initializer_list<type_t> allowed_types = {
-        type_t::float_type, type_t::half_type, type_t::fp8e4m3fnuz_type};
+    const std::initializer_list<type_t> allowed_types = {type_t::float_type,
+                                                         type_t::half_type,
+                                                         type_t::fp8e4m3fnuz_type,
+                                                         type_t::fp8e5m2fnuz_type,
+                                                         type_t::fp8e4m3fn_type,
+                                                         type_t::fp8e5m2_type};
     // Preliminary type check.
     if(not contains(allowed_types, result_type))
     {
@@ -702,6 +708,7 @@ struct find_mlir_standalone_op
                                     shape::type_t::half_type,
                                     shape::type_t::int8_type,
                                     shape::type_t::fp8e4m3fnuz_type,
+                                    shape::type_t::fp8e5m2fnuz_type,
                                     shape::type_t::fp8e4m3fn_type,
                                     shape::type_t::fp8e5m2_type},
                                    i->get_shape().type());

src/targets/gpu/gemm_impl.cpp

@@ -64,6 +64,7 @@ rocblas_datatype get_type(shape::type_t type)
     case shape::int32_type: return rocblas_datatype_i32_r;
     case shape::uint32_type: return rocblas_datatype_u32_r;
     case shape::fp8e4m3fnuz_type: return rocblas_datatype_f8_r;
+    case shape::fp8e5m2fnuz_type: return rocblas_datatype_bf8_r;
     case shape::fp8e4m3fn_type:
     case shape::fp8e5m2_type:
     case shape::tuple_type:

src/targets/gpu/hip_gemm_impl.cpp

@@ -70,6 +70,8 @@ hipDataType get_type_hipblas(shape::type_t type)
     case shape::int32_type: return HIP_R_32I;
     case shape::uint32_type: return HIP_R_32U;
     case shape::fp8e4m3fnuz_type: return HIP_R_8F_E4M3_FNUZ;
+    case shape::fp8e5m2fnuz_type:
+        return HIP_R_8F_E5M2_FNUZ;
 // TODO can remove this preprocessor conditional when hip verison defaults to have these types
 #ifdef ROCM_USE_FLOAT8
     case shape::fp8e4m3fn_type: return HIP_R_8F_E4M3;

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

@@ -37,6 +37,7 @@ namespace math {
 constexpr float as_float(migraphx::half x) { return x; }
 
 constexpr float as_float(migraphx::fp8::fp8e4m3fnuz x) { return x; }
+constexpr float as_float(migraphx::fp8::fp8e5m2fnuz x) { return x; }
 constexpr float as_float(migraphx::fp8::fp8e4m3fn x) { return x; }
 constexpr float as_float(migraphx::fp8::fp8e5m2 x) { return x; }
 
@@ -79,18 +80,22 @@ constexpr T as_float(T x)
         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...>())>                 \
-    auto __device__ name(migraphx::fp8::fp8e4m3fnuz x, Ts... xs) MIGRAPHX_RETURNS(     \
-        migraphx::fp8::fp8e4m3fnuz(fname(math::as_float(x), math::as_float(xs)...)))   \
-                                                                                       \
-        template <class... Ts, MIGRAPHX_REQUIRES(not is_any_vec<Ts...>())>             \
-        auto __device__ name(migraphx::fp8::fp8e4m3fn x, Ts... xs) MIGRAPHX_RETURNS(   \
-            migraphx::fp8::fp8e4m3fn(fname(math::as_float(x), math::as_float(xs)...))) \
-                                                                                       \
-            template <class... Ts, MIGRAPHX_REQUIRES(not is_any_vec<Ts...>())>         \
-            auto __device__ name(migraphx::fp8::fp8e5m2 x, Ts... xs) MIGRAPHX_RETURNS( \
-                migraphx::fp8::fp8e5m2(fname(math::as_float(x), math::as_float(xs)...)))
+#define MIGRAPHX_DEVICE_MATH_FP8(name, fname)                                              \
+    template <class... Ts, MIGRAPHX_REQUIRES(not is_any_vec<Ts...>())>                     \
+    auto __device__ name(migraphx::fp8::fp8e4m3fnuz x, Ts... xs) MIGRAPHX_RETURNS(         \
+        migraphx::fp8::fp8e4m3fnuz(fname(math::as_float(x), math::as_float(xs)...)))       \
+                                                                                           \
+        template <class... Ts, MIGRAPHX_REQUIRES(not is_any_vec<Ts...>())>                 \
+        auto __device__ name(migraphx::fp8::fp8e5m2fnuz x, Ts... xs) MIGRAPHX_RETURNS(     \
+            migraphx::fp8::fp8e5m2fnuz(fname(math::as_float(x), math::as_float(xs)...)))   \
+                                                                                           \
+            template <class... Ts, MIGRAPHX_REQUIRES(not is_any_vec<Ts...>())>             \
+            auto __device__ name(migraphx::fp8::fp8e4m3fn x, Ts... xs) MIGRAPHX_RETURNS(   \
+                migraphx::fp8::fp8e4m3fn(fname(math::as_float(x), math::as_float(xs)...))) \
+                                                                                           \
+                template <class... Ts, MIGRAPHX_REQUIRES(not is_any_vec<Ts...>())>         \
+                auto __device__ name(migraphx::fp8::fp8e5m2 x, Ts... xs) MIGRAPHX_RETURNS( \
+                    migraphx::fp8::fp8e5m2(fname(math::as_float(x), math::as_float(xs)...)))
 
 // Template with two overloads for math functions, one for half2 type and one for more generic
 // <half, N> vectorization where N is 4 or another even number.

src/targets/gpu/mlir.cpp

@@ -314,6 +314,8 @@ struct mlir_program
                 result = mlirF16TypeGet(ctx.get());
             else if(as.type_enum() == shape::fp8e4m3fnuz_type)
                 result = mlirFloat8E4M3FNUZTypeGet(ctx.get());
+            else if(as.type_enum() == shape::fp8e5m2fnuz_type)
+                result = mlirFloat8E5M2FNUZTypeGet(ctx.get());
             else if(as.type_enum() == shape::fp8e4m3fn_type)
                 result = mlirFloat8E4M3FNTypeGet(ctx.get());
             else if(as.type_enum() == shape::fp8e5m2_type)

src/targets/gpu/target.cpp

@@ -92,6 +92,7 @@ std::vector<pass> target::get_passes(migraphx::context& gctx, const compile_opti
     std::set<shape::type_t> unsupported_types(shape::types().begin(), shape::types().end());
     unsupported_types.erase(shape::type_t::float_type);
     unsupported_types.erase(shape::type_t::fp8e4m3fnuz_type);
+    unsupported_types.erase(shape::type_t::fp8e5m2fnuz_type);
     unsupported_types.erase(shape::type_t::fp8e4m3fn_type);
     unsupported_types.erase(shape::type_t::fp8e5m2_type);
     unsupported_types.erase(shape::type_t::half_type);
@@ -104,31 +105,31 @@ std::vector<pass> target::get_passes(migraphx::context& gctx, const compile_opti
     // whiltelist supported Ops for the FP8 types
     // different between fp8e4m3fnuz and OCP types because rocBLAS only has
     // support for fp8e4m3fnuz
-    std::set<std::string> unsupported_fp8e4m3fnuz_ops = {};
+    std::set<std::string> unsupported_fp8fnuz_ops = {};
     if(not gpu::rocblas_fp8_available())
     {
-        unsupported_fp8e4m3fnuz_ops.insert("dot");
-        unsupported_fp8e4m3fnuz_ops.insert("quant_dot");
+        unsupported_fp8fnuz_ops.insert("dot");
+        unsupported_fp8fnuz_ops.insert("quant_dot");
     }
 #if MIGRAPHX_USE_MIOPEN
     // MIOpen doesn't have support for fp8 pooling yet.
-    unsupported_fp8e4m3fnuz_ops.insert("pooling");
+    unsupported_fp8fnuz_ops.insert("pooling");
 #endif
     if(not gpu::gfx_has_fp8fnuz_intrinsics())
     {
-        unsupported_fp8e4m3fnuz_ops.insert("convolution");
-        unsupported_fp8e4m3fnuz_ops.insert("quant_convolution");
+        unsupported_fp8fnuz_ops.insert("convolution");
+        unsupported_fp8fnuz_ops.insert("quant_convolution");
     }
     // add all device kernels
-    unsupported_fp8e4m3fnuz_ops.insert("logsoftmax");
-    unsupported_fp8e4m3fnuz_ops.insert("nonzero");
-    unsupported_fp8e4m3fnuz_ops.insert("prefix_scan_sum");
-    unsupported_fp8e4m3fnuz_ops.insert("scatter_none");
-    unsupported_fp8e4m3fnuz_ops.insert("topk");
-    unsupported_fp8e4m3fnuz_ops.insert("rnn_var_sl_shift_output");
-    unsupported_fp8e4m3fnuz_ops.insert("multinomial");
-    unsupported_fp8e4m3fnuz_ops.insert("argmax");
-    unsupported_fp8e4m3fnuz_ops.insert("argmin");
+    unsupported_fp8fnuz_ops.insert("logsoftmax");
+    unsupported_fp8fnuz_ops.insert("nonzero");
+    unsupported_fp8fnuz_ops.insert("prefix_scan_sum");
+    unsupported_fp8fnuz_ops.insert("scatter_none");
+    unsupported_fp8fnuz_ops.insert("topk");
+    unsupported_fp8fnuz_ops.insert("rnn_var_sl_shift_output");
+    unsupported_fp8fnuz_ops.insert("multinomial");
+    unsupported_fp8fnuz_ops.insert("argmax");
+    unsupported_fp8fnuz_ops.insert("argmin");
 
     std::set<std::string> unsupported_fp8ocp_ops = {};
     // TODO: remove this when the flag is removed
@@ -193,7 +194,7 @@ std::vector<pass> target::get_passes(migraphx::context& gctx, const compile_opti
         dead_code_elimination{},
         prefuse_ops{},
         dead_code_elimination{},
-        eliminate_data_type{{migraphx::shape::fp8e4m3fnuz_type}, shape::float_type, unsupported_fp8e4m3fnuz_ops},
+        eliminate_data_type{{migraphx::shape::fp8e4m3fnuz_type, migraphx::shape::fp8e5m2fnuz_type}, shape::float_type, unsupported_fp8fnuz_ops},
         eliminate_data_type{{migraphx::shape::fp8e4m3fn_type, migraphx::shape::fp8e5m2_type}, shape::float_type, unsupported_fp8ocp_ops},
         dead_code_elimination{},
         rewrite_reduce{},

test/autocast_fp8.cpp

@@ -63,6 +63,7 @@ void autocast_fp8_1()
     EXPECT(m1 == m2);
 }
 TEST_CASE_REGISTER(autocast_fp8_1<migraphx::shape::fp8e4m3fnuz_type>);
+TEST_CASE_REGISTER(autocast_fp8_1<migraphx::shape::fp8e5m2fnuz_type>);
 TEST_CASE_REGISTER(autocast_fp8_1<migraphx::shape::fp8e4m3fn_type>);
 TEST_CASE_REGISTER(autocast_fp8_1<migraphx::shape::fp8e5m2_type>);
 
@@ -91,6 +92,7 @@ void autocast_fp8_2()
     EXPECT(m1 == m2);
 }
 TEST_CASE_REGISTER(autocast_fp8_2<migraphx::shape::fp8e4m3fnuz_type>);
+TEST_CASE_REGISTER(autocast_fp8_2<migraphx::shape::fp8e5m2fnuz_type>);
 TEST_CASE_REGISTER(autocast_fp8_2<migraphx::shape::fp8e4m3fn_type>);
 TEST_CASE_REGISTER(autocast_fp8_2<migraphx::shape::fp8e5m2_type>);
 
@@ -127,6 +129,7 @@ void autocast_fp8_3()
     EXPECT(m1 == m2);
 }
 TEST_CASE_REGISTER(autocast_fp8_3<migraphx::shape::fp8e4m3fnuz_type>);
+TEST_CASE_REGISTER(autocast_fp8_3<migraphx::shape::fp8e5m2fnuz_type>);
 TEST_CASE_REGISTER(autocast_fp8_3<migraphx::shape::fp8e4m3fn_type>);
 TEST_CASE_REGISTER(autocast_fp8_3<migraphx::shape::fp8e5m2_type>);
 
@@ -166,6 +169,7 @@ void autocast_fp8_4()
     EXPECT(m1 == m2);
 }
 TEST_CASE_REGISTER(autocast_fp8_4<migraphx::shape::fp8e4m3fnuz_type>);
+TEST_CASE_REGISTER(autocast_fp8_4<migraphx::shape::fp8e5m2fnuz_type>);
 TEST_CASE_REGISTER(autocast_fp8_4<migraphx::shape::fp8e4m3fn_type>);
 TEST_CASE_REGISTER(autocast_fp8_4<migraphx::shape::fp8e5m2_type>);