Add F4E2M1FN and F8E8M0FNU types

[sergey-kozub]

This PR adds F4E2M1FN primitive type (4-bit float with 2 bits exponent and 1 bit mantissa), F8E8M0FNU primitive type (8-bit float with 8 bits exponent, no mantissa and no sign) and enables loads/stores in the same way S4/U4 type is implemented.

This will enable using microscaling (MX) formats (RFC), such as MXFP4.

F4E2M1FN
- Exponent bias: 1
- Maximum stored exponent value: 3 (binary 11)
- Maximum unbiased exponent value: 3 - 1 = 2
- Minimum stored exponent value: 1 (binary 01)
- Minimum unbiased exponent value: 1 − 1 = 0
- Has Positive and Negative zero
- Doesn't have infinity
- Doesn't have NaNs

Additional details:
- Zeros (+/-): S.00.0
- Max normal number: S.11.1 = ±2^(2) x (1 + 0.5) = ±6.0
- Min normal number: S.01.0 = ±2^(0) = ±1.0
- Min subnormal number: S.00.1 = ±2^(0) x 0.5 = ±0.5

F8E8M0FNU
- Exponent bias: 127
- Maximum stored exponent value: 254 (binary 1111'1110)
- Maximum unbiased exponent value: 254 - 127 = 127
- Minimum stored exponent value: 0 (binary 0000'0000)
- Minimum unbiased exponent value: 0 − 127 = -127
- Doesn't have zero
- Doesn't have infinity
- NaN is encoded as binary 1111'1111

Additional details:
- Zeros cannot be represented
- Negative values cannot be represented
- Mantissa is always 1

Related PRs:

• Add MX floating point types (f4E2M1FN, f6E2M3FN, f6E3M2FN, f8E8M0FNU) stablehlo#2582
• Add sub-byte data types: float4_e2m1fn, float6_e2m3fn, float6_e3m2fn jax-ml/ml_dtypes#181
• [APFloat] Add APFloat support for FP4 data type llvm/llvm-project#95392
• [MLIR] Add f4E2M1FN type llvm/llvm-project#108877
• Add float8_e8m0_fnu (E8M0) OCP MX scale format jax-ml/ml_dtypes#166
• [APFloat] Add APFloat support for E8M0 type llvm/llvm-project#107127
• [MLIR] Add f8E8M0FNU type llvm/llvm-project#111028

The PR is split into multiple commits just to make the review easier, it is possible that some tests could fail if only some (i.e. not all) of these commits are applied.

[sergey-kozub]

Note: clang-format check fails on issues unrelated to this PR.

Also, running with clang-format 18.1.8 yields no issues.

[reedwm]

Sorry for the delay in reviewing.

You plan on adding the other MX types, right? If so, and convinient for you, could you add them in the same PR? Normally smaller PRs are better, but all these dtypes touch many of the same places in the code, so it's easier to batch review all of them at once. Granted, FP6 might be a bit weird if packed since their size doesn't divide the byte size -- those potentially should be in a separate PR.

But if adding the other MX types in inconvenient for you, happy to review this PR as-is.

[sergey-kozub]

You plan on adding the other MX types, right? If so, and convinient for you, could you add them in the same PR?

I'll add E8M0 implementation to this PR, as you suggest.
The FP6 types are not implemented yet, I'll look into that later (not a priority).

[sergey-kozub]

Added the F8E8M0FNU support patch to this PR. Will update the description soon after.

[reedwm]

Thanks for the PR!

Note I didn't review the MLIR codegen files (those in gpu/fusions/transforms) since I'm not familiar with those. @mooskagh can you find someone to review these files?

[pifon2a]

could you add to lower_tensors.mlir tests for the new 4-bit type similar to the existing tests for int4?

func.func @transfer_write_i4
func.func @transfer_read_i4
func.func @int4_constant

[sergey-kozub]

Added "f4_constant", "transfer_read_f4" and "transfer_write_f4".
BTW, there's no "transfer_read_i4" in that file.

[sergey-kozub]

I believe I addressed all the prior comments. PTAL when possible.

[pifon2a]

Thank you!

[mooskagh]

UPD: Disregard the comment, found it.

I'm trying to merge this PR.
It references float8.h in ml_dtypes.h, but I cannot find it either in this repository, nor internally (it was deleted a year ago).

[sergey-kozub]

Thanks, Reed! I'll proceed with rebasing 2533c35 to HEAD and create a new PR.

[reedwm]

I updated the ml_dtypes version used by TF in tensorflow/tensorflow@42bd1d3, and the TensorFlow Android build is still passing. So we should be able to merge the PR smoothly this time once you create the new PR :)

[sergey-kozub]

Created a new PR #21380 (rebased)