SWDEV-533546 - Add e8m0 conversions

hipamd/include/hip/amd_detail/amd_hip_fp8.h

@@ -65,6 +65,7 @@
 
 // Include it explicitly for HIPRTC
 #include "amd_hip_bf16.h"
+#include "amd_hip_mx_common.h"
 
 #if !defined(__HIPCC_RTC__)
 #include <hip/amd_detail/amd_hip_common.h>
@@ -949,6 +950,124 @@ __FP8_HOST_STATIC__ __hip_fp8x2_storage_t __hip_cvt_halfraw2_to_fp8x2(
   return __hip_cvt_float2_to_fp8x2(__half22float2(__half2(x)), sat, interp);
 }
 
+namespace hip_detail {
+
+constexpr __hip_fp8_storage_t e8m0_NaN = 0xFFU;
+constexpr std::uint16_t bf16_NaN = 0x7FFFU;
+
+constexpr std::uint16_t bf16_sig_mask = 0x007FU;
+constexpr std::uint32_t float_sig_mask = 0x007FFFFFU;
+constexpr std::uint64_t double_sig_mask = 0x000FFFFFFFFFFFFFU;
+
+constexpr std::uint16_t bf16_max_exp = 0x7F80U;
+constexpr std::uint32_t float_max_exp = 0x7F800000U;
+constexpr std::uint64_t double_max_exp = 0x7FF0000000000000U;
+
+constexpr std::uint16_t bf16_sign_mask = 0x8000U;
+constexpr std::uint32_t float_sign_mask = 0x80000000U;
+constexpr std::uint64_t double_sign_mask = 0x8000000000000000U;
+
+constexpr std::uint16_t bf16_half_sig_bit = 0x0040U;
+constexpr std::uint32_t float_half_sig_bit = 0x00400000U;
+constexpr std::uint64_t double_half_sig_bit = 0x0008000000000000U;
+
+}  // namespace hip_detail
+
+__FP8_HOST_DEVICE_STATIC__ __hip_fp8_storage_t __hip_cvt_double_to_e8m0(
+    const double val, const __hip_saturation_t saturate, const enum hipRoundMode rounding) {
+  union {
+    double as_double;
+    std::uint64_t as_int;
+  } u{val};
+
+  // Shifts out mantisa bits from double dtype
+  unsigned short double_exp =
+      static_cast<unsigned short>((~hip_detail::double_sign_mask & u.as_int) >> 52);
+  __hip_fp8_storage_t e8m0;
+  if (double_exp == 0x0U) {
+    e8m0 = 0x0U;
+  } else if ((double_exp - 0x0380U) > 0x00FF) {  // if double is NaN/Inf/or too large
+    e8m0 = hip_detail::e8m0_NaN;
+  } else {
+    e8m0 =
+        double_exp - 0x0380U;  // shift due to bias difference between double and single precision
+  }
+
+  // If there is a mantisa and the exp wont overflow round up
+  if ((rounding == hipRoundPosInf) && (u.as_int & hip_detail::double_sig_mask) &&
+      (!((u.as_int & ~hip_detail::double_sign_mask) < hip_detail::double_half_sig_bit)) &&
+      (e8m0 < hip_detail::e8m0_NaN)) {
+    ++e8m0;
+  }
+
+  // If e8m0 is NaN and exponent is a large non-inf value round down to a value
+  if ((saturate == __HIP_SATFINITE) && (e8m0 == hip_detail::e8m0_NaN) &&
+      ((u.as_int & ~hip_detail::double_sign_mask) <= hip_detail::double_max_exp)) {
+    --e8m0;
+  }
+  return e8m0;
+}
+
+__FP8_HOST_DEVICE_STATIC__ __hip_fp8_storage_t __hip_cvt_float_to_e8m0(
+    const float val, const __hip_saturation_t saturate, const enum hipRoundMode rounding) {
+  union {
+    float as_float;
+    std::uint32_t as_int;
+  } u{val};
+
+  // Shifts out mantisa bits from float dtype
+  __hip_fp8_storage_t e8m0 = static_cast<unsigned char>(u.as_int >> 23);
+
+  // If there is a mantisa and the exp wont overflow round up
+  if ((rounding == hipRoundPosInf) && (u.as_int & hip_detail::float_sig_mask) &&
+      (!((u.as_int & ~hip_detail::float_sign_mask) < hip_detail::float_half_sig_bit)) &&
+      (e8m0 < hip_detail::e8m0_NaN)) {
+    ++e8m0;
+  }
+
+  // If e8m0 is NaN and exponent is a large non-inf value round down to a value
+  if ((saturate == __HIP_SATFINITE) && (e8m0 == hip_detail::e8m0_NaN) &&
+      ((u.as_int & ~hip_detail::float_sign_mask) <= hip_detail::float_max_exp)) {
+    --e8m0;
+  }
+  return e8m0;
+}
+
+__FP8_HOST_DEVICE_STATIC__ __hip_fp8_storage_t
+__hip_cvt_bfloat16raw_to_e8m0(const __hip_bfloat16_raw hr, const __hip_saturation_t saturate,
+                              const enum hipRoundMode rounding) {
+  // Shifts out mantisa bits from bf16 dtype
+  __hip_fp8_storage_t e8m0 = static_cast<unsigned char>(hr.x >> 7);
+
+  // If there is a mantisa and the exp wont overflow round up
+  if ((rounding == hipRoundPosInf) && (hr.x & hip_detail::bf16_sig_mask) &&
+      (!((hr.x & ~hip_detail::bf16_sign_mask) < hip_detail::bf16_half_sig_bit)) &&
+      (e8m0 < hip_detail::e8m0_NaN)) {
+    ++e8m0;
+  }
+
+  // If e8m0 is NaN and exponent is a large non-inf value round down to a value
+  if ((saturate == __HIP_SATFINITE) && (e8m0 == hip_detail::e8m0_NaN) &&
+      ((hr.x & ~hip_detail::bf16_sign_mask) <= hip_detail::bf16_max_exp)) {
+    --e8m0;
+  }
+  return e8m0;
+}
+
+__FP8_HOST_DEVICE_STATIC__ __hip_bfloat16_raw
+__hip_cvt_e8m0_to_bf16raw(const __hip_fp8_storage_t x) {
+  switch (x) {
+    case 0x00U:
+      return __hip_bfloat16_raw{0x0040U};
+      break;
+    case hip_detail::e8m0_NaN:
+      return __hip_bfloat16_raw{hip_detail::bf16_NaN};
+      break;
+    default:
+      return __hip_bfloat16_raw{static_cast<unsigned short>(x << 7)};
+  }
+}
+
 /**
  * \brief struct representing single fp8 number with e4m3 interpretation
  *