Skip to content

Commit d68ab07

Browse files
apivovarovapivovarov
authored
[RFC] Add f8E4M3 and f8E3M4 types support (#2486)
### Summary This is a proposal to add `Float8E4M3` and `Float8E3M4` floating point types to StableHLO. Feedback welcome, see [RFC: Float8E4M3 and Float8E3M4](https://github.com/apivovarov/stablehlo/blob/rfc_f8E4M3_f8E3M4/rfcs/20240808-f8E4M3_f8E3M4.md) for more details. ### References and Links - LLVM [PR-97179](llvm/llvm-project#97179) [APFloat] Add support for f8E4M3 IEEE 754 type (Merged) - LLVM [PR-97118](llvm/llvm-project#97118) [MLIR] Add f8E4M3 IEEE 754 type (Merged) - LLVM [PR-99698](llvm/llvm-project#99698) [APFloat] Add support for f8E3M4 IEEE 754 type (Merged) - LLVM [PR-101230](llvm/llvm-project#101230) [MLIR] Add f8E3M4 IEEE 754 type (Merged) - [RFC: FP8 in StableHLO](https://github.com/openxla/stablehlo/blob/main/rfcs/20221031-fp8.md) - [RFC: Float8E4M3FNUZ and Float8E5M2FNUZ](https://github.com/openxla/stablehlo/blob/main/rfcs/20230321-fp8_fnuz.md) - StableHLO [PR-2482](#2482) Add f8E4M3 and f8E3M4 types support - [Amazon EC2 Trn1 Instances](https://aws.amazon.com/ec2/instance-types/trn1/) - ml_dtypes [PR-161](jax-ml/ml_dtypes#161) Add float8_e4m3 (Merged) - ml_dtypes [PR-171](jax-ml/ml_dtypes#171) Add float8_e3m4 (Merged) - XLA [PR-16585](openxla/xla#16585) Add support for float8_e4m3
1 parent 1456dfa commit d68ab07

File tree

1 file changed

+178
-0
lines changed

1 file changed

+178
-0
lines changed

‎rfcs/20240808-f8E4M3_f8E3M4.md

Lines changed: 178 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,178 @@
1+
# RFC: Float8E4M3 and Float8E3M4
2+
3+
Status: In Review<br/>
4+
Initial version: 8/8/2024<br/>
5+
Last updated: 8/9/2024<br/>
6+
Discussion thread: [PR-2486](https://github.com/openxla/stablehlo/pull/2486)
7+
[RFC] Add f8E4M3 and f8E3M4 types support
8+
9+
## Summary
10+
11+
Amazon has proposed two new FP8 types, Float8E4M3 and Float8E3M4. These
12+
types are implemented in commercially available hardware[^1], and added to MLIR
13+
builtin types[^2]˒[^3] and LLVM APFloat[^4]˒[^5].
14+
15+
Both Float8E4M3 and Float8E3M4 follows IEEE 754 convention similar to existing
16+
type Float8E5M2.
17+
18+
### Float8E4M3
19+
20+
8-bit floating point type with 1 sign bit, 4 bits exponent and 3 bits mantissa
21+
following IEEE-754 conventions with bit layout S1E4M3.
22+
23+
```c
24+
f8E4M3 (IEEE 754)
25+
- Exponent bias: 7
26+
- Minimum stored exponent value: 1 (binary 0001)
27+
- Maximum stored exponent value: 14 (binary 1110)
28+
- Minimum unbiased exponent value: 1 − 7 = −6
29+
- Maximum unbiased exponent value: 14 - 7 = 7
30+
- Precision specifies the total number of bits used for the significand
31+
(mantisa), including implicit leading integer bit = 3 + 1 = 4
32+
- Follows IEEE 754 conventions for representation of special values
33+
- Has Positive and Negative zero
34+
- Has Positive and Negative infinity
35+
- Has NaNs
36+
37+
Additional details:
38+
- Min exp (unbiased): -6
39+
- Max exp (unbiased): 7
40+
- Infinities (+/-): S.1111.000
41+
- Zeros (+/-): S.0000.000
42+
- NaNs: S.1111.{001, 010, 011, 100, 101, 110, 111}
43+
- Min normal number: S.0001.000 = +/-2^(1 - 7) x (1 + 0) = +/-2^(-6)
44+
- Max normal number: S.1110.111 = +/-2^(14 - 7) x (1 + 7/8) = +/-240
45+
- Min subnormal number: S.0000.001 = +/-2^(-6) x 1/8 = +/-2^(-9)
46+
- Max subnormal number: S.0000.111 = +/-2^(-6) x 7/8 = +/-2^(-9) x 7
47+
```
48+
49+
#### Comparison of Float8E4M3FN and Float8E4M3
50+
51+
| |Float8E4M3FN |Float8E4M3 |
52+
|-------------------|------------------------------------------------------------------------|-------------------------------------------------------------------------|
53+
|Bias |7 |7 |
54+
|Min Normal Value |`0bS0001000` = -1<sup>S</sup> $\times$ 1.0 $\times$ 2<sup>-6</sup> |`0bS0001000` = -1<sup>S</sup> $\times$ 1.0 $\times$ 2<sup>-6</sup> |
55+
|Max Normal Value |`0bS1111110` = -1<sup>S</sup> $\times$ 1.75 $\times$ 2<sup>8</sup> = 448|`0bS1110111` = -1<sup>S</sup> $\times$ 1.875 $\times$ 2<sup>7</sup> = 240|
56+
|Min Subnormal Value|`0bS0000001` = -1<sup>S</sup> $\times$ 0.125 $\times$ 2<sup>-6</sup> |`0bS0000001` = -1<sup>S</sup> $\times$ 0.125 $\times$ 2<sup>-6</sup> |
57+
|Max Subnormal Value|`0bS0000111` = -1<sup>S</sup> $\times$ 0.875 $\times$ 2<sup>-6</sup> |`0bS0000111` = -1<sup>S</sup> $\times$ 0.875 $\times$ 2<sup>-6</sup> |
58+
|NaN |`0bS1111111` |`0bS1111MMM`, where `MMM` is non-zero. |
59+
|Infinity |N/A |`0bS1111000` |
60+
|-0.0 |`0b10000000` |`0b10000000` |
61+
62+
### Float8E3M4
63+
64+
8-bit floating point type with 1 sign bit, 3 bits exponent and 4 bits mantissa
65+
following IEEE-754 conventions with bit layout S1E3M4.
66+
67+
```c
68+
f8E3M4 (IEEE 754)
69+
- Exponent bias: 3
70+
- Minimum stored exponent value: 1 (binary 001)
71+
- Maximum stored exponent value: 6 (binary 110)
72+
- Minimum unbiased exponent value: 1 − 3 = −2
73+
- Maximum unbiased exponent value: 6 - 3 = 3
74+
- Precision specifies the total number of bits used for the significand
75+
(mantissa), including implicit leading integer bit = 4 + 1 = 5
76+
- Follows IEEE 754 conventions for representation of special values
77+
- Has Positive and Negative zero
78+
- Has Positive and Negative infinity
79+
- Has NaNs
80+
81+
Additional details:
82+
- Min exp (unbiased): -2
83+
- Max exp (unbiased): 3
84+
- Infinities (+/-): S.111.0000
85+
- Zeros (+/-): S.000.0000
86+
- NaNs: S.111.{0,1}⁴ except S.111.0000
87+
- Min normal number: S.001.0000 = +/-2^(1 - 3) x (1 + 0) = +/-0.25
88+
- Max normal number: S.110.1111 = +/-2^(6 - 3) x (1 + 15/16) = +/-15.5
89+
- Min subnormal number: S.000.0001 = +/-2^(-2) x 1/16 = +/-2^(-6)
90+
- Max subnormal number: S.000.1111 = +/-2^(-2) x 15/16 = +/-2^(-6) x 15
91+
```
92+
93+
### Comparison of Float8E5M2, Float8E4M3 and Float8E3M4
94+
95+
| |Float8E5M2 |Float8E4M3 |Float8E3M4 |
96+
|-------------------|----------------------------------------------------------------------------|-------------------------------------------------------------------------|---------------------------------------------------------------------------|
97+
|Bias |15 |7 |3 |
98+
|Min Normal Value |`0bS0000100` = -1<sup>S</sup> $\times$ 1.0 $\times$ 2<sup>-14</sup> |`0bS0001000` = -1<sup>S</sup> $\times$ 1.0 $\times$ 2<sup>-6</sup> |`0bS0010000` = -1<sup>S</sup> $\times$ 1.0 $\times$ 2<sup>-2</sup> |
99+
|Max Normal Value |`0bS1111011` = -1<sup>S</sup> $\times$ 1.75 $\times$ 2<sup>15</sup> = 57344 |`0bS1110111` = -1<sup>S</sup> $\times$ 1.875 $\times$ 2<sup>7</sup> = 240|`0bS1101111` = -1<sup>S</sup> $\times$ 1.9375 $\times$ 2<sup>3</sup> = 15.5|
100+
|Min Subnormal Value|`0bS0000001` = -1<sup>S</sup> $\times$ 0.25 $\times$ 2<sup>-14</sup> |`0bS0000001` = -1<sup>S</sup> $\times$ 0.125 $\times$ 2<sup>-6</sup> |`0bS0000001` = -1<sup>S</sup> $\times$ 0.0625 $\times$ 2<sup>-2</sup> |
101+
|Max Subnormal Value|`0bS0000011` = -1<sup>S</sup> $\times$ 0.75 $\times$ 2<sup>-14</sup> |`0bS0000111` = -1<sup>S</sup> $\times$ 0.875 $\times$ 2<sup>-6</sup> |`0bS0001111` = -1<sup>S</sup> $\times$ 0.9375 $\times$ 2<sup>-2</sup> |
102+
|NaN |`0bS11111MM`, where `MM` is non-zero. |`0bS1111MMM`, where `MMM` is non-zero. |`0bS111MMMM`, where `MMMM` is non-zero. |
103+
|Infinity |`0bS1111100` |`0bS1111000` |`0bS1110000` |
104+
|-0.0 |`0b10000000` |`0b10000000` |`0b10000000` |
105+
106+
## Changes in StableHLO
107+
108+
I propose adding Float8E4M3 and Float8E3M4 types to StableHLO similar to the
109+
previously introduces FP8 types (below) with some differences:
110+
111+
- [FP8 RFC](https://github.com/openxla/xla/discussions/22)
112+
- [[RFC] Add Float8E4M3FNUZ and Float8E5M2FNUZ to StableHLO](https://github.com/openxla/stablehlo/pull/1342)
113+
114+
### StableHLO Interpreter
115+
116+
To provide a reference implementation, I intend to add support for
117+
Float8E4M3 and Float8E3M4 in the StableHLO interpreter. This will be
118+
useful for testing other backends and validating new implementations. This will
119+
be achieved in two ways:
120+
121+
1. Map directly to the appropriate APFloat operation.
122+
2. Cast up to the appropriate type, use that implementation, cast back down.
123+
124+
### Float8E4M3 and Float8E3M4 Arithmetic
125+
126+
I intend for Float8E4M3 and Float8E3M4 to be types that support the
127+
appropriate arithmetic operations, like any other floating point type. For
128+
platforms that don't have hardware support for these types, they may either
129+
throw an error and reject the program or cast up to an appropriate higher
130+
precision type that is supported, compute the answer, and cast back down.
131+
132+
This is a simple approach that aligns with user expectations of a floating
133+
point data type, and is the approach taken by BFloat16. This also gives
134+
backends freedom to exploit any hardware support.
135+
136+
Here's an example of a real JAX program (logging the MLIR) computing a simple
137+
dot product in Float8E4M3. Note the answer is slightly "wrong", as expected
138+
due to the lower precision (round-to-nearest).
139+
140+
```python
141+
>>> import jax
142+
>>> import jax.numpy as jnp
143+
>>> x = jnp.arange(8, dtype=jnp.float8_e4m3)
144+
module @jit_iota {
145+
func.func public @main() -> tensor<8xf8E4M3> {
146+
%0 = stablehlo.iota dim = 0 : tensor<8xf8E4M3>
147+
return %0 : tensor<8xf8E4M3>
148+
}
149+
}
150+
>>> x
151+
Array([0, 1, 2, 3, 4, 5, 6, 7], dtype=float8_e4m3)
152+
>>> x @ x
153+
module @jit_matmul {
154+
func.func public @main(%arg0: tensor<8xf8E4M3> {mhlo.sharding = ""}, %arg1: tensor<8xf8E4M3> {mhlo.sharding = ""}) -> tensor<f8E4M3> {
155+
%0 = "stablehlo.dot_general"(%arg0, %arg1) {dot_dimension_numbers = #stablehlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]} : (tensor<8xf8E4M3>, tensor<8xf8E4M3>) -> tensor<f8E4M3>
156+
return %0 : tensor<f8E4M3>
157+
}
158+
}
159+
Array(144, dtype=float8_e4m3)
160+
```
161+
162+
### Testing
163+
164+
Built on the StableHLO interpreter, I intend to introduce tests for all
165+
possible operations with Float8E4M3 and Float8E3M4 inputs. This will at
166+
a minimum mean adding additional cases to the `interpret_X.mlir` family of
167+
tests.
168+
169+
### References and Links
170+
171+
- [RFC: FP8 in StableHLO](https://github.com/openxla/stablehlo/blob/main/rfcs/20221031-fp8.md)
172+
- [RFC: Float8E4M3FNUZ and Float8E5M2FNUZ](https://github.com/openxla/stablehlo/blob/main/rfcs/20230321-fp8_fnuz.md)
173+
174+
[^1]: [Amazon EC2 Trn1 Instances](https://aws.amazon.com/ec2/instance-types/trn1/)
175+
[^2]: LLVM [PR-97118](https://github.com/llvm/llvm-project/pull/97118) [MLIR] Add f8E4M3 IEEE 754 type (Merged)
176+
[^3]: LLVM [PR-101230](https://github.com/llvm/llvm-project/pull/101230) [MLIR] Add f8E3M4 IEEE 754 type (Merged)
177+
[^4]: LLVM [PR-97179](https://github.com/llvm/llvm-project/pull/97179) [APFloat] Add support for f8E4M3 IEEE 754 type (Merged)
178+
[^5]: LLVM [PR-99698](https://github.com/llvm/llvm-project/pull/99698) [APFloat] Add support for f8E3M4 IEEE 754 type (Merged)

0 commit comments

Comments
 (0)