1 #pragma once
2
3 /// Defines the Float8_e4m3fnuz 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 Float8_e4m3fn:
8 /// s eeee mmm
9 /// 1 sign bit
10 /// 4 exponent bits
11 /// 3 mantissa bits
12 /// The key differences versus Float8_e4m3fn are:
13 /// bias = 8
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/Export.h>
21 #include <c10/macros/Macros.h>
22 #include <c10/util/floating_point_utils.h>
23 #include <type_traits>
24
25 #if defined(__cplusplus)
26 #include <cstdint>
27 #elif !defined(__OPENCL_VERSION__)
28 #include <math.h>
29 #include <stdint.h>
30 #endif
31
32 #include <iosfwd>
33 #include <ostream>
34
35 namespace c10 {
36
37 namespace detail {
38
39 /*
40 * Convert a 32-bit floating-point number in IEEE single-precision format to a
41 * 8-bit floating-point number in fp8 E4M3FNUZ format, in bit representation.
42 */
fp8e4m3fnuz_from_fp32_value(float f)43 inline C10_HOST_DEVICE uint8_t fp8e4m3fnuz_from_fp32_value(float f) {
44 /*
45 * Binary representation of 256.0f, which is the first value not representable
46 * (i.e. the first value which would overflow in to the sign bit, resulting in
47 * a NaN) in fp8e4m3fnuz range:
48 * 1 0000 000 - fp8e4m3fnuz
49 * 0 10000111 00000000000000000000000 - fp32
50 */
51 constexpr uint32_t fnuz_max = UINT32_C(0x87) << 23;
52
53 /*
54 * A mask for converting fp32 numbers lower than fp8e4m3fnuz normal range
55 * into denorm representation
56 * magic number: ((127 - 8) + (23 - 3) + 1)
57 */
58 constexpr uint32_t denorm_mask = UINT32_C(0x8C) << 23;
59
60 uint32_t f_bits = fp32_to_bits(f);
61
62 uint32_t result = 0u;
63
64 /*
65 * Extract the sign of the input number into the high bit of the 32-bit word:
66 *
67 * +---+----------------------------------+
68 * | S |0000000 00000000 00000000 00000000|
69 * +---+----------------------------------+
70 * Bits 31 0-31
71 */
72 const uint32_t sign = f_bits & UINT32_C(0x80000000);
73
74 /*
75 * Set sign bit to 0
76 */
77 f_bits ^= sign;
78
79 if (f_bits >= fnuz_max) {
80 // NaN -- sign bit set to 1, rest 0s.
81 return 0x80;
82 }
83
84 if (f_bits < (UINT32_C(0x78) << 23) /* 2^-7 in float32 */) {
85 // Input exponent is less than -7, the smallest e4m3fnuz exponent, so the
86 // number will become subnormal.
87 f_bits = fp32_to_bits(fp32_from_bits(f_bits) + fp32_from_bits(denorm_mask));
88 result = static_cast<uint8_t>(f_bits - denorm_mask);
89 if (result == 0) {
90 // fnuz types don't have negative zero.
91 return 0;
92 }
93 } else {
94 // resulting mantissa is odd
95 uint8_t mant_odd = (f_bits >> 20) & 1;
96
97 // update exponent, rounding bias part 1
98 f_bits += ((uint32_t)(8 - 127) << 23) + 0x7FFFF;
99
100 // rounding bias part 2
101 f_bits += mant_odd;
102
103 // take the bits!
104 result = static_cast<uint8_t>(f_bits >> 20);
105 }
106
107 result |= sign >> 24;
108 return result;
109 }
110
111 } // namespace detail
112
113 struct alignas(1) Float8_e4m3fnuz {
114 uint8_t x;
115
116 struct from_bits_t {};
from_bitsFloat8_e4m3fnuz117 C10_HOST_DEVICE static constexpr from_bits_t from_bits() {
118 return from_bits_t();
119 }
120
121 Float8_e4m3fnuz() = default;
122
Float8_e4m3fnuzFloat8_e4m3fnuz123 constexpr C10_HOST_DEVICE Float8_e4m3fnuz(uint8_t bits, from_bits_t)
124 : x(bits) {}
125 inline C10_HOST_DEVICE Float8_e4m3fnuz(float value);
126 inline C10_HOST_DEVICE operator float() const;
127 inline C10_HOST_DEVICE bool isnan() const;
128 };
129
130 C10_API inline std::ostream& operator<<(
131 std::ostream& out,
132 const Float8_e4m3fnuz& value) {
133 out << (float)value;
134 return out;
135 }
136
137 } // namespace c10
138
139 #include <c10/util/Float8_e4m3fnuz-inl.h> // IWYU pragma: keep
140