-
Notifications
You must be signed in to change notification settings - Fork 3k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[Feature] Add bfloat16 support for CPU
- Loading branch information
Showing
13 changed files
with
297 additions
and
35 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,68 @@ | ||
| /** | ||
| * Copyright (c) 2023 by Contributors | ||
| * @file dgl/runtime/ndarray.h | ||
| * @brief BFloat16 CPU header | ||
| */ | ||
| #ifndef DGL_RUNTIME_BFLOAT16_H_ | ||
| #define DGL_RUNTIME_BFLOAT16_H_ | ||
|
|
||
| #include <cmath> | ||
|
|
||
| class BFloat16 { | ||
| uint16_t val; | ||
|
|
||
| public: | ||
| constexpr BFloat16() : val(0) {} | ||
| // disabling "explicit" warning in lint | ||
| // implicit usage on costructor is expected | ||
| BFloat16(float f) { // NOLINT | ||
| if (std::isnan(f)) { | ||
| val = 0x7FC0; | ||
| } else { | ||
| union { | ||
| uint16_t iraw16[2]; | ||
| uint32_t iraw32; | ||
| float f32; | ||
| }; | ||
|
|
||
| f32 = f; | ||
| const uint32_t rounding_bias = 0x00007FFF + (iraw16[1] & 0x1); | ||
| val = static_cast<uint16_t>((iraw32 + rounding_bias) >> 16); | ||
| } | ||
| } | ||
| static constexpr BFloat16 Min() { | ||
| BFloat16 min; | ||
| min.val = 0xFF80; | ||
| return min; | ||
| } | ||
|
|
||
| static constexpr BFloat16 Max() { | ||
| BFloat16 max; | ||
| max.val = 0x7F80; | ||
| return max; | ||
| } | ||
|
|
||
| BFloat16& operator-=(const float& rhs) { | ||
| float lhs = (*this); | ||
| (*this) = lhs - rhs; | ||
| return *this; | ||
| } | ||
|
|
||
| BFloat16& operator+=(const float& rhs) { | ||
| float lhs = (*this); | ||
| (*this) = lhs + rhs; | ||
| return *this; | ||
| } | ||
|
|
||
| operator float() const { | ||
| union { | ||
| float f; | ||
| uint16_t raw[2]; | ||
| }; | ||
| raw[0] = 0; | ||
| raw[1] = val; | ||
| return f; | ||
| } | ||
| }; | ||
|
|
||
| #endif // DGL_RUNTIME_BFLOAT16_H_ |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.