1 #pragma once
2
3 /// Defines the Float8_e5m2fnuz type (8-bit floating-point) including
4 /// conversions to standard C types and basic arithmetic operations. Note that
5 /// arithmetic operations are implemented by converting to floating point and
6 /// performing the operation in float32.
7 /// Binary configuration remains the same as e5m2:
8 /// s eeeee mm
9 /// 1 sign bit
10 /// 5 exponent bits
11 /// 2 mantissa bits
12 /// The key differences that e5m2fnuz brings are:
13 /// bias = 16
14 /// no infinities or negative zero
15 /// NaN only when sign bit is 1, rest all 0s
16 ///
17 /// Implementation based on the paper https://arxiv.org/pdf/2206.02915.pdf and
18 /// the existing Float8_e4m3fn implementation.
19
20 #include <c10/macros/Macros.h>
21 #include <c10/util/TypeSafeSignMath.h>
22 #include <c10/util/floating_point_utils.h>
23
24 #if defined(__cplusplus)
25 #include <cstdint>
26 #elif !defined(__OPENCL_VERSION__)
27 #include <math.h>
28 #include <stdint.h>
29 #endif
30
31 #include <iosfwd>
32 #include <ostream>
33
34 namespace c10 {
35
36 namespace detail {
37
38 /*
39 * Convert a 32-bit floating-point number in IEEE single-precision format to a
40 * 8-bit floating-point number in fp8 E5M2 format, in bit representation.
41 */
fp8e5m2fnuz_from_fp32_value(float f)42 inline C10_HOST_DEVICE uint8_t fp8e5m2fnuz_from_fp32_value(float f) {
43 /*
44 * Binary representation of 65536.0f, which is the first value not
45 * representable (i.e. the first value which would overflow in to the sign
46 * bit, resulting in a NaN) in fp8e4m3fnuz range:
47 * 1 00000 00 - fp8e5m2fnuz
48 * 0 10001111 00000000000000000000000 - fp32
49 */
50 constexpr uint32_t fnuz_max = UINT32_C(0x8F) << 23;
51
52 /*
53 * A mask for converting fp32 numbers lower than fp8e5m2fnuz normal range
54 * into denormalized representation.
55 * magic number: ((127 - 16) + (23 - 2) + 1)
56 */
57 constexpr uint32_t denorm_mask = UINT32_C(0x85) << 23;
58
59 uint32_t f_bits = fp32_to_bits(f);
60 uint32_t result = 0u;
61
62 /*
63 * Extract the sign of the input number into the high bit of the 32-bit word:
64 *
65 * +---+----------------------------------+
66 * | S |0000000 00000000 00000000 00000000|
67 * +---+----------------------------------+
68 * Bits 31 0-31
69 */
70 const uint32_t sign = f_bits & UINT32_C(0x80000000);
71
72 /*
73 * Set sign bit to 0
74 */
75 f_bits ^= sign;
76
77 if (f_bits >= fnuz_max) {
78 // NaN -- sign bit set to 1, rest 0s
79 return 0x80;
80 }
81
82 if (f_bits < (UINT32_C(0x70) << 23) /* 2^-15 in float32 */) {
83 // Input exponent is less than -15, the smallest e5m2fnuz exponent, so the
84 // number will become subnormal.
85 f_bits = fp32_to_bits(fp32_from_bits(f_bits) + fp32_from_bits(denorm_mask));
86 result = static_cast<uint8_t>(f_bits - denorm_mask);
87 if (result == 0) {
88 // fnuz types don't have negative zero.
89 return 0;
90 }
91 } else {
92 // resulting mantissa is odd
93 uint8_t mant_odd = (f_bits >> 21) & 1;
94
95 // update exponent, rounding bias part 1
96 f_bits += ((uint32_t)(16 - 127) << 23) + 0xFFFFF;
97
98 // rounding bias part 2
99 f_bits += mant_odd;
100
101 // take the bits!
102 result = static_cast<uint8_t>(f_bits >> 21);
103 }
104
105 result |= sign >> 24;
106 return result;
107 }
108
109 } // namespace detail
110
111 struct alignas(1) Float8_e5m2fnuz {
112 uint8_t x;
113
114 struct from_bits_t {};
from_bitsFloat8_e5m2fnuz115 C10_HOST_DEVICE static constexpr from_bits_t from_bits() {
116 return from_bits_t();
117 }
118
119 Float8_e5m2fnuz() = default;
120
Float8_e5m2fnuzFloat8_e5m2fnuz121 constexpr C10_HOST_DEVICE Float8_e5m2fnuz(uint8_t bits, from_bits_t)
122 : x(bits) {}
123 inline C10_HOST_DEVICE Float8_e5m2fnuz(float value);
124 inline C10_HOST_DEVICE operator float() const;
125 inline C10_HOST_DEVICE bool isnan() const;
126 inline C10_HOST_DEVICE bool isinf() const;
127 };
128
129 C10_API inline std::ostream& operator<<(
130 std::ostream& out,
131 const Float8_e5m2fnuz& value) {
132 out << (float)value;
133 return out;
134 }
135
136 } // namespace c10
137
138 #include <c10/util/Float8_e5m2fnuz-inl.h> // IWYU pragma: keep
139