Add support for float8_e4m3 and float8_e3m4 types

[apivovarov]

This PR adds f8E4M3 and f8E3M4 types support to XLA (mainly to cpu_compiler).

f8E4M3 type follows IEEE 754 convention.

f8E4M3 (IEEE 754)
- Exponent bias: 7
- Maximum stored exponent value: 14 (binary 1110)
- Maximum unbiased exponent value: 14 - 7 = 7
- Minimum stored exponent value: 1 (binary 0001)
- Minimum unbiased exponent value: 1 − 7 = −6
- Precision specifies the total number of bits used for the significand (mantisa), 
    including implicit leading integer bit = 3 + 1 = 4
- Follows IEEE 754 conventions for representation of special values
- Has Positive and Negative zero
- Has Positive and Negative infinity
- Has NaNs

Additional details:
- Max exp (unbiased): 7
- Min exp (unbiased): -6
- Infinities (+/-): S.1111.000
- Zeros (+/-): S.0000.000
- NaNs: S.1111.{001, 010, 011, 100, 101, 110, 111}
- Max normal number: S.1110.111 = +/-2^(7) x (1 + 0.875) = +/-240
- Min normal number: S.0001.000 = +/-2^(-6)
- Max subnormal number: S.0000.111 = +/-2^(-6) x 0.875 = +/-2^(-9) x 7
- Min subnormal number: S.0000.001 = +/-2^(-6) x 0.125 = +/-2^(-9)

f8E3M4 type follows IEEE 754 convention

f8E3M4 (IEEE 754)
- Exponent bias: 3
- Maximum stored exponent value: 6 (binary 110)
- Maximum unbiased exponent value: 6 - 3 = 3
- Minimum stored exponent value: 1 (binary 001)
- Minimum unbiased exponent value: 1 − 3 = −2
- Precision specifies the total number of bits used for the significand (mantissa), 
    including implicit leading integer bit = 4 + 1 = 5
- Follows IEEE 754 conventions for representation of special values
- Has Positive and Negative zero
- Has Positive and Negative infinity
- Has NaNs

Additional details:
- Max exp (unbiased): 3
- Min exp (unbiased): -2
- Infinities (+/-): S.111.0000
- Zeros (+/-): S.000.0000
- NaNs: S.111.{0,1}⁴ except S.111.0000
- Max normal number: S.110.1111 = +/-2^(6-3) x (1 + 15/16) = +/-2^3 x 31 x 2^(-4) = +/-15.5
- Min normal number: S.001.0000 = +/-2^(1-3) x (1 + 0) = +/-2^(-2)
- Max subnormal number: S.000.1111 = +/-2^(-2) x 15/16 = +/-2^(-2) x 15 x 2^(-4) = +/-15 x 2^(-6)
- Min subnormal number: S.000.0001 = +/-2^(-2) x 1/16 =  +/-2^(-2) x 2^(-4) = +/-2^(-6)

Testing:

bazel test \
//xla:array2d_test \
//xla:fp_util_test \
//xla:literal_comparison_test \
//xla:literal_test \
//xla/mlir/utils:type_util_test \
//xla:primitive_util_test \
//xla/python/ifrt:dtype_test \
//xla/python:xla_client_test \
//xla/service:elemental_ir_emitter_test \
//xla/service:float_normalization_test \
//xla/service/gpu/tests:float_conversions_test \
//xla/tests:array_elementwise_ops_test \
//xla/tests:constants_test \
//xla/tests:convert_test \
//xla/tests:float8_test \
//xla:util_test

bazel test \
//xla/hlo/translate/hlo_to_mhlo/tests:import.hlo.test \
//xla/hlo/translate/mhlo_to_hlo/tests:export.mlir.test \
//xla/mlir_hlo/tests:Dialect/mhlo/hlo-legalize-to-stablehlo.mlir.test \
//xla/mlir_hlo/tests:Dialect/mhlo/ops.mlir.test \
//xla/mlir_hlo/tests:Dialect/mhlo/stablehlo-legalize-to-hlo.mlir.test

Related PRs:

• LLVM PR-97179 [APFloat] Add support for f8E4M3 IEEE 754 type (Merged)
• LLVM PR-97118 [MLIR] Add f8E4M3 IEEE 754 type (Merged)
• LLVM PR-99698 [APFloat] Add support for f8E3M4 IEEE 754 type (Merged)
• LLVM PR-101230 [MLIR] Add f8E3M4 IEEE 754 type (Merged)
• StableHLO PR-2486 [RFC] Add f8E4M3 and f8E3M4 types support (Merged)
• StableHLO PR-2482 Add f8E4M3 and f8E3M4 types support (Merged)
• ml_dtypes PR-161 Add float8_e4m3 (Merged)
• ml_dtypes PR-171 Add float8_e3m4 (Merged)
• XLA PR-17075 [TSL] Bump ml_dtypes. Add float8_e4m3, float8_e3m4 (Approved)
• XLA PR-3200 Add support for float8_e4m3fnuz and float8_e5m2fnuz (Template)
• JAX PR-23585 Add float8_e4m3 type support (in Review)

[mooskagh]

Probably not intended to fetch it from your repo.

[mooskagh]

If it's intended that you also update the deps in the same PR, could you split it in separate PRs?

[apivovarov]

This PR is pending the merge of StableHLO openxla/stablehlo#2482 Add f8E4M3 and f8E3M4 types support (in Review).

[apivovarov]

Hi Alexander, Integrate StableHLO at openxla/stablehlo@4f31b2e7 was werged to XLA main today. It includes float8_e4m3 type support. My temporary change in third_party/stablehlo/workspace.bzl was removed from this PR. @mooskagh

[apivovarov]

Reed, David, could you please help review this PR? @reedwm @ddunl

[ddunl]

I think Reed is the best person to review, I think this will require a patch on our end due to tensorflow/third_party, let me know when you approve and I can take care of the patch

[apivovarov]

Hi Reed,

This PR introduces support for the new f8E4M3 type, which adheres to the IEEE-754 convention. I've already added this type to the LLVM, MLIR, ml_dtypes, and StableHLO projects. This PR extends the support to XLA and includes a reference implementation for the CPU compiler.

Could you please help review this PR?

@reedwm

[reedwm]

Thanks for the change! Sorry for the delay in reviewing.

Normally, it's better to minimize the size of PRs, but I would prefer if E3M4 is also added in the same PR, since it touches most of the same files in the exact same way as E4M3, so it makes it easy to batch review both dtypes at once.

But if adding E3M4 to the same PR is inconvenient with you, I'm fine with this being done as a separate, future PR.

[reedwm]

Since the type list here is growing, I would avoid hardcoding the list of FP8 types. One way to avoid this is to have DoWithUpcastToF32 take a should_upcast bool instead of the existing upcast_types list. Then you can pass something like should_upcast = BitWidth(b.GetShape(x).element_type) <= 16.

There are a lot of places where we list out all FP8 types, but every place we can remove listing these out will help when more types are added :)

[apivovarov]

Opened PR #17130 - Add default upcasting behavior to DoWithUpcastToF32

[reedwm]

Since this test is almost identical to F8E4M3FNDistance, can you merge them to avoid duplication?

One way to to create a type-parameterized test with TYPED_TEST_P. Another way would be to have a for-loop over primtiives types F8E4M3 and F8E4M3FN, and in the body use primitive_util::PrimitiveTypeSwitch with a lambda that does the CalculateDistanceInFloats calls. See here for an example of how to use PrimitiveTypeSwitch.

[apivovarov]

Opened PR #17135 - Add TypeParam to FP8E4M3DistanceTest

[reedwm]

Don't change the formatting here.

[apivovarov]

restored

[apivovarov]

github workflow runs pipx to check clang formatting. Opened PR #17234 Format hlo_evaluator_typed_visitor.h

[reedwm]

I think all the FP8 tests are duplicated for each FP8 type. Can you use TYPED_TEST_P to deduplicate them?

[apivovarov]

Opened #17133 - Dedup LiteralComparisonTests

[reedwm]

Not sure why there is a convention of using subsequent single letters to name the FP8 values (q16, r16, s16, etc) but you should follow it or change the convention. Either name this q16, renaming the above e5m2 to p16, or rename all the other FP8 variable names to something more descriptive, as you did for this one.

[apivovarov]

renamed to q16

[reedwm]

To avoid having to modify this every time a new FP8 type is added, remove all these FP8 cases and check if IsF8Type(literal.shape().element_type() before the switch statement.

[apivovarov]

Opened PR #17170 Code dedup in execution_trace_utils LiteralToValue

[jreiffers]

Just FYI, this was an intentional choice. Missing switch cases are a compiler error, so having a switch without a default case is preferable when possible. No big deal though.

[apivovarov]

Using single-case switch statements can make it easier for the compiler to detect potential errors in the code.
Opened PR #17279 - Use switch case without default in LiteralToValue

[reedwm]

Merge this with the existing test ResolveIfUnsupportedF8e5m2, either by looping over values (F8E4M3, F8E5M2) or by using a value-parameterized test with TEST_P.

[apivovarov]

Opened PR #17177 - Parametrize FloatNormalizationF8Test ResolveIfUnsupportedF8

[reedwm]

This comment is actually no longer true. We should change the other tests with this comment as well. The test OneCellF8e5m2fnuz does have an FP8 output, so you can use that as an example in modifying this test.

If you want, you can also change the two existing tests with the comment to have FP8 outputs as well.

[apivovarov]

Opened PR #17182 - Parametrize ConstantsFloatTest OneCellFloat

[reedwm]

This is effectively subtracting 8 from the exponent I think, as the difference in exponent bias is 8. Why not do that subtraction directly?

[apivovarov]

In contrast to the EmitF16ToF8e4m3fn function, the EmitF16ToF8e4m3 function does not include special code to handle Inf and NaN cases. (code around constexpr int max_finite_value = 0x5F7F;)

If I use -8 approach then several tests in //xla/tests:convert_test_cpu FAILED.
e.g.

• inf -> -1.0
• nan -> -1.5

Example:

input is inf
EmitReducePrecisionIR returns
x = 0.11111.0000000000 (0x7C00)

Option1: minus 8
x -= 0.01000.0000000000
// x is 0.10111.0000000000
// Shift to convert to F8: x is 1.0111.000
// f8e4m3 Result is -1.0 (Wrong)

Option2: Right shift E5 exponent's leftmost bit
x = (x & 0b1001'1111'1111'1111) | ((x & 0b0100'0000'0000'0000) >> 1)
// x is 0.01111.0000000000
// Shift to convert to F8: x is 0.1111.000
// f8e4m3 Result is inf (Correct)

[reedwm]

0.11111.000000000000 has 12 zeros at the end, when it should have 10.

[apivovarov]

fixed

[reedwm]

Also, I think you might have to modify expand_float_ops.cc, which is used by the new MLIR emitters which replace the existing emitters on GPUs. But I'm not very familiar with these new emitters. @jreiffers, can you advice on what needs to be done here

[jreiffers]

We have generic support for all possible float conversions, but the emitted code might not be optimal, so it should be considered a fallback. I didn't look at these conversion routines here in detail, but if they're better, it would make sense to port them to the MLIR pipeline.

[apivovarov]

I updated xla/service/gpu/fusions/transforms/expand_float_ops.cc and added f8E4M3 cases to:

• IsInf()
• IsNaN()
• RewriteF8Cst::matchAndRewrite() // If we're comparing to +-0, compare the absolute values.

expand_float_ops.cc includes a specialized function for the f8e5m2 type - EmitF16ToF8e5m2(). This is because F16 is technically f16e5m10. The two types are similar, with the primary difference being that the mantissa in f8e5m2 is truncated to 2 bits.

f8E4M3 has a different number of exponent and mantissa bits. The conversion can be efficiently managed using the "generic support for all possible float conversions".

Tested xla on CUDA:

//xla/tests/...    799 tests: 799 tests pass
//xla/service/...  865 tests: 865 tests pass
//xla/client/...    77 tests: 77 tests pass
//xla/runtime/...    1 test: 1 test passes
//xla/ffi/...        6 tests: 6 tests pass
//xla/hlo/...       12 tests: 12 tests pass
//xla/mlir/...     141 tests: 141 tests pass
//xla/mlir_hlo/...  98 tests: 98 tests pass
//xla/pjrt/...      26 tests: 26 tests pass
//xla/tools/...     41 tests: 41 tests pass
//xla/translate/... 61 tests: 61 tests pass

[jreiffers]

Apologies for the delay, I'm OOO this week. Will take a look on Monday.

[apivovarov]

PR has been merged! Reed, thank you for your help, guidance, and support! @reedwm

[reedwm]

No problem, and thanks for the well-tested PR! Also thank you for all the test clean up PRs!

Note in merging, when converting to E3M4, I had to change the code to first convert to half to take into account we do not use an ml-dtypes version that includes jax-ml/ml_dtypes#205 yet. I added TODOs in the form of TODO(b/370786669) to all places where a conversion to half was added. These can be removed once we update ml-dtypes to a version that includes jax-ml/ml_dtypes#205. I'm happy to do this myself or for you to do this once ml-dtypes is updated.