Add support for float8_e4m3fnuz and float8_e5m2fnuz. by jakeh-gc · Pull Request #3200 · openxla/xla

jakeh-gc · 2023-05-25T16:07:13Z

This adds support for the two FP8 types float8_e4m3fnuz and float8_e5m2fnuz to XLA similar to float8_e4m3fn, float8_e4m3b11, and float8_e5m2.

jakeh-gc · 2023-05-25T16:26:27Z

I see where that test is failing... I'll get that fixed.

reedwm · 2023-05-25T18:26:02Z

@cantonios can you review the TSL changes (I'll also take a look at them). For some reason, I cannot assign you as a reviewer to this PR.

@burmako is there anything we need to do on the StableHLO side before merging this?

cantonios · 2023-05-25T18:46:07Z

@cantonios can you review the TSL changes (I'll also take a look at them). For some reason, I cannot assign you as a reviewer to this PR.

Probably because I'm not part of the openxla team. But yeah, happy to.

cantonios

This change will likely need to wait until TensorFlow/TSL switches over to use ml_dtypes (as described in another comment). I'm in the process of doing this... my best estimate is within the next week or so. At that point, none of the TSL changes here will be necessary.

burmako · 2023-05-25T19:41:08Z

@reedwm Thank you for reaching out! These types have gone through the StableHLO RFC process and are now part of the StableHLO spec, so I don't think anything further is needed on the StableHLO side.

reedwm

Haven't reviewed elemental_ir_emitter.cc yet, I'll try to get to that tomorrow.

burmako

LGTM for MHLO/HLO parity

reedwm

Please add tests to convert_test.cc and constants_test.cc similar to existing F8 tests in those files.

Also note I didn't review TSL changes based on @cantonios's comments that the dependency on TSL float types would go away.

reedwm

Please add tests to convert_test.cc and constants_test.cc similar to existing F8 tests in those files.

Also note I didn't review TSL changes based on @cantonios's comments that the dependency on TSL float types would go away.

Imported from GitHub PR openxla/xla#3200 This adds support for the two FP8 types `float8_e4m3fnuz` and `float8_e5m2fnuz` to XLA similar to `float8_e4m3fn`, `float8_e4m3b11`, and `float8_e5m2`. Copybara import of the project: -- 3b96f8fe219c1ac1bec5c4b99ff9c51148706981 by Jake Hall <[email protected]>: Add support for float8_e4m3fnuz and float8_e5m2fnuz. Merging this change closes #3200 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#3200 from jakeh-gc:fp8_fnuz 3b96f8fe219c1ac1bec5c4b99ff9c51148706981 PiperOrigin-RevId: 543802274

reedwm · 2023-06-27T22:33:12Z

@jakeh-gc, I'm still working on merging this. Please don't commit to the PR in the meantime, as it's hard to update the internal changes with the PR changes.

Imported from GitHub PR openxla/xla#3200 This adds support for the two FP8 types `float8_e4m3fnuz` and `float8_e5m2fnuz` to XLA similar to `float8_e4m3fn`, `float8_e4m3b11`, and `float8_e5m2`. Copybara import of the project: -- 3b96f8fe219c1ac1bec5c4b99ff9c51148706981 by Jake Hall <[email protected]>: Add support for float8_e4m3fnuz and float8_e5m2fnuz. Merging this change closes #3200 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#3200 from jakeh-gc:fp8_fnuz 3b96f8fe219c1ac1bec5c4b99ff9c51148706981 PiperOrigin-RevId: 543802274

Imported from GitHub PR openxla/xla#3200 This adds support for the two FP8 types `float8_e4m3fnuz` and `float8_e5m2fnuz` to XLA similar to `float8_e4m3fn`, `float8_e4m3b11`, and `float8_e5m2`. Copybara import of the project: -- 3b96f8fe219c1ac1bec5c4b99ff9c51148706981 by Jake Hall <[email protected]>: Add support for float8_e4m3fnuz and float8_e5m2fnuz. Merging this change closes #3200 PiperOrigin-RevId: 544198797

FUTURE_COPYBARA_INTEGRATE_REVIEW=#3200 from jakeh-gc:fp8_fnuz 3b96f8f PiperOrigin-RevId: 544197768

Imported from GitHub PR openxla/xla#16585 This PR adds f8E4M3 and f8E3M4 types support to XLA (mainly to cpu_compiler). ### `f8E4M3` type follows IEEE 754 convention. ```c 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 ```c 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](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) - StableHLO [PR-2486](openxla/stablehlo#2486) [RFC] Add f8E4M3 and f8E3M4 types support (Merged) - StableHLO [PR-2482](openxla/stablehlo#2482) Add f8E4M3 and f8E3M4 types support (Merged) - 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-17075](openxla/xla#17075) [TSL] Bump ml_dtypes. Add float8_e4m3, float8_e3m4 (Approved) - XLA [PR-3200](openxla/xla#3200) Add support for float8_e4m3fnuz and float8_e5m2fnuz (Template) - JAX [PR-23585](jax-ml/jax#23585) Add float8_e4m3 type support (in Review) Copybara import of the project: -- ec1c723027012a816d7e17f268c5f034863696e6 by Alexander Pivovarov <[email protected]>: Add support for float8_e4m3 and float8_e3m4 types Merging this change closes #16585 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#16585 from apivovarov:float8_e4m3 ec1c723027012a816d7e17f268c5f034863696e6 PiperOrigin-RevId: 680651037

github-actions Bot added the kokoro:force-run Forces CI to rerun label May 25, 2023

github-actions Bot assigned xla-rotation May 25, 2023

kokoro-team removed the kokoro:force-run Forces CI to rerun label May 25, 2023

jakeh-gc force-pushed the fp8_fnuz branch from 222b60f to 5612eb6 Compare May 25, 2023 16:30

github-actions Bot added the kokoro:force-run Forces CI to rerun label May 25, 2023

kokoro-team removed the kokoro:force-run Forces CI to rerun label May 25, 2023

cheshire assigned reedwm May 25, 2023

reedwm self-requested a review May 25, 2023 18:23

cantonios suggested changes May 25, 2023

View reviewed changes

Comment thread third_party/tsl/tsl/platform/float8.h Outdated

Comment thread third_party/tsl/tsl/python/lib/core/custom_casts.cc Outdated

Comment thread third_party/tsl/tsl/python/lib/core/float8.cc Outdated

Comment thread third_party/tsl/tsl/python/lib/core/float8.h Outdated

burmako suggested changes May 25, 2023

View reviewed changes

Comment thread xla/translate/hlo_to_mhlo/hlo_utils.cc Outdated

Comment thread xla/translate/mhlo_to_hlo/type_to_shape.cc Outdated

jakeh-gc force-pushed the fp8_fnuz branch from 5612eb6 to 59d7098 Compare May 25, 2023 23:27

github-actions Bot added the kokoro:force-run Forces CI to rerun label May 25, 2023

kokoro-team removed the kokoro:force-run Forces CI to rerun label May 25, 2023

reedwm suggested changes May 26, 2023

View reviewed changes

Comment thread xla/primitive_util.h Outdated

Comment thread xla/primitive_util_test.cc Outdated

Comment thread xla/util_test.cc Outdated

Comment thread xla/xla_data.proto Outdated

jakeh-gc force-pushed the fp8_fnuz branch from 59d7098 to 547910c Compare May 26, 2023 09:05

github-actions Bot added the kokoro:force-run Forces CI to rerun label May 26, 2023

kokoro-team removed the kokoro:force-run Forces CI to rerun label May 26, 2023

jakeh-gc force-pushed the fp8_fnuz branch from 547910c to faf80e4 Compare May 26, 2023 13:53

github-actions Bot added the kokoro:force-run Forces CI to rerun label May 26, 2023

kokoro-team removed the kokoro:force-run Forces CI to rerun label May 26, 2023

burmako approved these changes May 26, 2023

View reviewed changes

reedwm suggested changes Jun 1, 2023

View reviewed changes

jakeh-gc force-pushed the fp8_fnuz branch from faf80e4 to 906e7bd Compare June 15, 2023 17:22

github-actions Bot added the kokoro:force-run Forces CI to rerun label Jun 15, 2023

kokoro-team removed the kokoro:force-run Forces CI to rerun label Jun 15, 2023

copybara-service Bot mentioned this pull request Jun 27, 2023

PR #3200: Add support for float8_e4m3fnuz and float8_e5m2fnuz. #3859

Closed

copybara-service Bot mentioned this pull request Jun 27, 2023

PR #3200: Add support for float8_e4m3fnuz and float8_e5m2fnuz. google/tsl#835

Closed

copybara-service Bot mentioned this pull request Jun 29, 2023

[PJRT] Add buffer type getter to PJRT C and C++ APIs. #3879

Closed

copybara-service Bot closed this in 59a27d1 Jun 29, 2023

copybara-service Bot mentioned this pull request Jun 29, 2023

[PJRT C API] Add tests for PJRT C API Implementation through a test factory #3444

Closed

copybara-service Bot pushed a commit that referenced this pull request Jun 29, 2023

Testing CHECK macro

18dc4dc

FUTURE_COPYBARA_INTEGRATE_REVIEW=#3200 from jakeh-gc:fp8_fnuz 3b96f8f PiperOrigin-RevId: 544197768

copybara-service Bot mentioned this pull request Jun 29, 2023

Testing CHECK macro #3912

Closed

apivovarov mentioned this pull request Sep 11, 2024

Add support for float8_e4m3 and float8_e3m4 types #16585

Closed

This was referenced Sep 30, 2024

PR #16585: Add support for float8_e4m3 and float8_e3m4 types google/tsl#2762

Merged

PR #16585: Add support for float8_e4m3 and float8_e3m4 types #17774

Merged

PR #16585: Add support for float8_e4m3 and float8_e3m4 types tensorflow/tensorflow#76821

Merged

ScXfjiang mentioned this pull request Oct 6, 2024

[ROCM] Add NANOO FP8 Support in TensorFlow tensorflow/tensorflow#77117

Closed

Conversation

jakeh-gc commented May 25, 2023

Uh oh!

jakeh-gc commented May 25, 2023

Uh oh!

reedwm commented May 25, 2023

Uh oh!

cantonios commented May 25, 2023

Uh oh!

cantonios left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

burmako commented May 25, 2023

Uh oh!

Uh oh!

Uh oh!

reedwm left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

burmako left a comment

Choose a reason for hiding this comment

Uh oh!

reedwm left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

reedwm left a comment

Choose a reason for hiding this comment

Uh oh!

reedwm commented Jun 27, 2023

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants