1 #pragma once
2
3 /// Defines the Float8_e5m2 type (8-bit floating-point) including conversions
4 /// to standard C types and basic arithmetic operations. Note that arithmetic
5 /// operations are implemented by converting to floating point and
6 /// performing the operation in float32.
7 /// Binary configuration:
8 /// s eeeee mm
9 /// 1 sign bit
10 /// 5 exponent bits
11 /// 2 mantissa bits
12 /// bias = 15
13 ///
14 /// Implementation based on the paper https://arxiv.org/pdf/2209.05433.pdf
15 /// and inspired by Half implementation from pytorch/c10/util/Half.h
16
17 #include <c10/util/Half.h>
18
19 namespace c10 {
20
21 namespace detail {
22
23 /*
24 * Convert a 8-bit floating-point number in fp8 E5M2 format, in bit
25 * representation, to a 32-bit floating-point number in IEEE single-precision
26 * format, in bit representation.
27 *
28 * @note The implementation doesn't use any floating-point operations.
29 */
fp8e5m2_to_fp32_value(uint8_t input)30 inline C10_HOST_DEVICE float fp8e5m2_to_fp32_value(uint8_t input) {
31 /*
32 * Extend the fp8 E5M2 number to 32 bits and shift to the
33 * upper part of the 32-bit word:
34 * +---+----+---+-----------------------------+
35 * | S |EEEEE|MM|0000 0000 0000 0000 0000 0000|
36 * +---+----+---+-----------------------------+
37 * Bits 31 26-30 24-25 0-23
38 *
39 * S - sign bit, E - bits of the biased exponent, M - bits of the mantissa, 0
40 * - zero bits.
41 */
42 uint16_t half_representation = input;
43 half_representation <<= 8;
44 return fp16_ieee_to_fp32_value(half_representation);
45 }
46
47 /*
48 * Convert a 32-bit floating-point number in IEEE single-precision format to a
49 * 8-bit floating-point number in fp8 E5M2 format, in bit representation.
50 */
fp8e5m2_from_fp32_value(float f)51 inline C10_HOST_DEVICE uint8_t fp8e5m2_from_fp32_value(float f) {
52 /*
53 * Binary representation of fp32 infinity
54 * 0 11111111 00000000000000000000000
55 */
56 constexpr uint32_t fp32_inf = UINT32_C(255) << 23;
57
58 /*
59 * Binary representation of 65536.0f, which is the first value
60 * not representable in fp8e5m2 range:
61 * 0 11111 00 - fp8e5m2
62 * 0 10001111 00000000000000000000000 - fp32
63 */
64 constexpr uint32_t fp8_max = UINT32_C(143) << 23;
65
66 /*
67 * A mask for converting fp32 numbers lower than fp8e5m2 normal range
68 * into denorm representation
69 * magic number: ((127 - 15) + (23 - 2) + 1)
70 */
71 constexpr uint32_t denorm_mask = UINT32_C(134) << 23;
72
73 uint32_t f_bits = fp32_to_bits(f);
74 uint8_t result = 0u;
75
76 /*
77 * Extract the sign of the input number into the high bit of the 32-bit word:
78 *
79 * +---+----------------------------------+
80 * | S |0000000 00000000 00000000 00000000|
81 * +---+----------------------------------+
82 * Bits 31 0-31
83 */
84 const uint32_t sign = f_bits & UINT32_C(0x80000000);
85
86 /*
87 * Set sign bit to 0
88 */
89 f_bits ^= sign;
90
91 if (f_bits >= fp8_max) {
92 // NaN - all exponent and mantissa bits set to 1
93 result = f_bits > fp32_inf ? UINT8_C(0x7F) : UINT8_C(0x7C);
94 } else {
95 if (f_bits < (UINT32_C(113) << 23)) {
96 // Input number is smaller than 2^(-14), which is the smallest
97 // fp8e5m2 normal number
98 f_bits =
99 fp32_to_bits(fp32_from_bits(f_bits) + fp32_from_bits(denorm_mask));
100 result = static_cast<uint8_t>(f_bits - denorm_mask);
101 } else {
102 // resulting mantissa is odd
103 uint32_t mant_odd = (f_bits >> 21) & 1;
104
105 // update exponent, rounding bias part 1
106 f_bits += ((uint32_t)(15 - 127) << 23) + 0xFFFFF;
107
108 // rounding bias part 2
109 f_bits += mant_odd;
110
111 // take the bits!
112 result = static_cast<uint8_t>(f_bits >> 21);
113 }
114 }
115
116 result |= static_cast<uint8_t>(sign >> 24);
117 return result;
118 }
119
120 } // namespace detail
121
122 struct alignas(1) Float8_e5m2 {
123 uint8_t x;
124
125 struct from_bits_t {};
from_bitsFloat8_e5m2126 C10_HOST_DEVICE static constexpr from_bits_t from_bits() {
127 return from_bits_t();
128 }
129
130 Float8_e5m2() = default;
131
Float8_e5m2Float8_e5m2132 constexpr C10_HOST_DEVICE Float8_e5m2(uint8_t bits, from_bits_t) : x(bits) {}
133 inline C10_HOST_DEVICE Float8_e5m2(float value);
134 inline C10_HOST_DEVICE operator float() const;
135 inline C10_HOST_DEVICE bool isnan() const;
136 inline C10_HOST_DEVICE bool isinf() const;
137 };
138
139 C10_API inline std::ostream& operator<<(
140 std::ostream& out,
141 const Float8_e5m2& value) {
142 out << (float)value;
143 return out;
144 }
145
146 } // namespace c10
147
148 #include <c10/util/Float8_e5m2-inl.h> // IWYU pragma: keep
149