1 | pub(crate) fn f32_to_bf16(value: f32) -> u16 { |
---|---|

2 | // Convert to raw bytes |

3 | let x = value.to_bits(); |

4 | |

5 | // check for NaN |

6 | if x & 0x7FFF_FFFFu32 > 0x7F80_0000u32 { |

7 | // Keep high part of current mantissa but also set most significiant mantissa bit |

8 | return ((x >> 16) | 0x0040u32) as u16; |

9 | } |

10 | |

11 | // round and shift |

12 | let round_bit = 0x0000_8000u32; |

13 | if (x & round_bit) != 0 && (x & (3 * round_bit - 1)) != 0 { |

14 | (x >> 16) as u16 + 1 |

15 | } else { |

16 | (x >> 16) as u16 |

17 | } |

18 | } |

19 | |

20 | pub(crate) fn f64_to_bf16(value: f64) -> u16 { |

21 | // Convert to raw bytes, truncating the last 32-bits of mantissa; that precision will always |

22 | // be lost on half-precision. |

23 | let val = value.to_bits(); |

24 | let x = (val >> 32) as u32; |

25 | |

26 | // Extract IEEE754 components |

27 | let sign = x & 0x8000_0000u32; |

28 | let exp = x & 0x7FF0_0000u32; |

29 | let man = x & 0x000F_FFFFu32; |

30 | |

31 | // Check for all exponent bits being set, which is Infinity or NaN |

32 | if exp == 0x7FF0_0000u32 { |

33 | // Set mantissa MSB for NaN (and also keep shifted mantissa bits). |

34 | // We also have to check the last 32 bits. |

35 | let nan_bit = if man == 0 && (val as u32 == 0) { |

36 | 0 |

37 | } else { |

38 | 0x0040u32 |

39 | }; |

40 | return ((sign >> 16) | 0x7F80u32 | nan_bit | (man >> 13)) as u16; |

41 | } |

42 | |

43 | // The number is normalized, start assembling half precision version |

44 | let half_sign = sign >> 16; |

45 | // Unbias the exponent, then bias for bfloat16 precision |

46 | let unbiased_exp = ((exp >> 20) as i64) - 1023; |

47 | let half_exp = unbiased_exp + 127; |

48 | |

49 | // Check for exponent overflow, return +infinity |

50 | if half_exp >= 0xFF { |

51 | return (half_sign | 0x7F80u32) as u16; |

52 | } |

53 | |

54 | // Check for underflow |

55 | if half_exp <= 0 { |

56 | // Check mantissa for what we can do |

57 | if 7 - half_exp > 21 { |

58 | // No rounding possibility, so this is a full underflow, return signed zero |

59 | return half_sign as u16; |

60 | } |

61 | // Don't forget about hidden leading mantissa bit when assembling mantissa |

62 | let man = man | 0x0010_0000u32; |

63 | let mut half_man = man >> (14 - half_exp); |

64 | // Check for rounding |

65 | let round_bit = 1 << (13 - half_exp); |

66 | if (man & round_bit) != 0 && (man & (3 * round_bit - 1)) != 0 { |

67 | half_man += 1; |

68 | } |

69 | // No exponent for subnormals |

70 | return (half_sign | half_man) as u16; |

71 | } |

72 | |

73 | // Rebias the exponent |

74 | let half_exp = (half_exp as u32) << 7; |

75 | let half_man = man >> 13; |

76 | // Check for rounding |

77 | let round_bit = 0x0000_1000u32; |

78 | if (man & round_bit) != 0 && (man & (3 * round_bit - 1)) != 0 { |

79 | // Round it |

80 | ((half_sign | half_exp | half_man) + 1) as u16 |

81 | } else { |

82 | (half_sign | half_exp | half_man) as u16 |

83 | } |

84 | } |

85 | |

86 | pub(crate) fn bf16_to_f32(i: u16) -> f32 { |

87 | // If NaN, keep current mantissa but also set most significiant mantissa bit |

88 | if i & 0x7FFFu16 > 0x7F80u16 { |

89 | f32::from_bits((i as u32 | 0x0040u32) << 16) |

90 | } else { |

91 | f32::from_bits((i as u32) << 16) |

92 | } |

93 | } |

94 | |

95 | pub(crate) fn bf16_to_f64(i: u16) -> f64 { |

96 | // Check for signed zero |

97 | if i & 0x7FFFu16 == 0 { |

98 | return f64::from_bits((i as u64) << 48); |

99 | } |

100 | |

101 | let half_sign = (i & 0x8000u16) as u64; |

102 | let half_exp = (i & 0x7F80u16) as u64; |

103 | let half_man = (i & 0x007Fu16) as u64; |

104 | |

105 | // Check for an infinity or NaN when all exponent bits set |

106 | if half_exp == 0x7F80u64 { |

107 | // Check for signed infinity if mantissa is zero |

108 | if half_man == 0 { |

109 | return f64::from_bits((half_sign << 48) | 0x7FF0_0000_0000_0000u64); |

110 | } else { |

111 | // NaN, keep current mantissa but also set most significiant mantissa bit |

112 | return f64::from_bits((half_sign << 48) | 0x7FF8_0000_0000_0000u64 | (half_man << 45)); |

113 | } |

114 | } |

115 | |

116 | // Calculate double-precision components with adjusted exponent |

117 | let sign = half_sign << 48; |

118 | // Unbias exponent |

119 | let unbiased_exp = ((half_exp as i64) >> 7) - 127; |

120 | |

121 | // Check for subnormals, which will be normalized by adjusting exponent |

122 | if half_exp == 0 { |

123 | // Calculate how much to adjust the exponent by |

124 | let e = (half_man as u16).leading_zeros() - 9; |

125 | |

126 | // Rebias and adjust exponent |

127 | let exp = ((1023 - 127 - e) as u64) << 52; |

128 | let man = (half_man << (46 + e)) & 0xF_FFFF_FFFF_FFFFu64; |

129 | return f64::from_bits(sign | exp | man); |

130 | } |

131 | // Rebias exponent for a normalized normal |

132 | let exp = ((unbiased_exp + 1023) as u64) << 52; |

133 | let man = (half_man & 0x007Fu64) << 45; |

134 | f64::from_bits(sign | exp | man) |

135 | } |

136 |