Coverage for src / monte_neo / core / acceleration / float8.py: 20%

1""" 

2Float8 precision support for GPU acceleration. 

3 

4Supports E4M3 and E5M2 formats for memory bandwidth optimization. 

5""" 

6 

7from __future__ import annotations 

8 

9import numpy as np 

10 

11 

12class Float8Encoder: 

13 """Encoder for float8 formats.""" 

14 

15 def __init__(self, format_type: str = "e4m3") -> None: 

16 """ 

17 Initialize float8 encoder. 

18  

19 Args: 

20 format_type: 'e4m3' or 'e5m2' 

21 """ 

22 self.format_type = format_type 

23 if format_type == "e4m3": 

24 self.mantissa_bits = 3 

25 self.exponent_bits = 4 

26 self.bias = 7 

27 self.max_exp = 15 

28 elif format_type == "e5m2": 

29 self.mantissa_bits = 2 

30 self.exponent_bits = 5 

31 self.bias = 15 

32 self.max_exp = 31 

33 else: 

34 raise ValueError(f"Unsupported format: {format_type}") 

35 

36 self.mantissa_mask = (1 << self.mantissa_bits) - 1 

37 self.exponent_mask = (1 << self.exponent_bits) - 1 

38 

39 def encode_scalar(self, value: float) -> int: 

40 """Encode single float32 to float8.""" 

41 if np.isnan(value): 

42 return 0x7F if self.format_type == "e4m3" else 0x7C 

43 

44 if value == 0.0: 

45 return 0 

46 

47 # Handle infinity 

48 if np.isinf(value): 

49 sign = 1 if value < 0 else 0 

50 exp = self.exponent_mask 

51 mant = 0 

52 return (sign << 7) | (exp << self.mantissa_bits) | mant 

53 

54 # Extract sign 

55 sign = 1 if value < 0 else 0 

56 value = abs(value) 

57 

58 # Convert to binary representation 

59 bits = np.float32(value).view(np.uint32) 

60 

61 # Extract exponent and mantissa from float32 

62 float32_bias = 127 

63 # Cast to int to avoid unsigned subtraction issues 

64 float32_exp = int((bits >> 23) & 0xFF) - float32_bias 

65 float32_mant = int(bits & 0x7FFFFF) 

66 

67 # Convert exponent 

68 mantissa_shift = 23 - self.mantissa_bits 

69 if float32_exp >= self.max_exp - self.bias: 

70 # Overflow -> infinity 

71 exp = self.exponent_mask 

72 mant = 0 

73 elif float32_exp < -self.bias: 

74 # Underflow -> zero 

75 return 0 

76 else: 

77 exp = float32_exp + self.bias 

78 # Convert mantissa 

79 # Shift mantissa to match float8 precision 

80 mant = (float32_mant >> mantissa_shift) & self.mantissa_mask 

81 

82 # Round if necessary (round to nearest even) 

83 if mantissa_shift > 0: 

84 round_bit = (float32_mant >> (mantissa_shift - 1)) & 1 

85 sticky_bits = float32_mant & ((1 << (mantissa_shift - 1)) - 1) 

86 if round_bit and (sticky_bits or (mant & 1)): 

87 mant += 1 

88 if mant > self.mantissa_mask: 

89 mant = 0 

90 exp += 1 

91 if exp > self.exponent_mask: 

92 exp = self.exponent_mask 

93 

94 return (sign << 7) | (exp << self.mantissa_bits) | mant 

95 

96 def decode_scalar(self, encoded: int) -> float: 

97 """Decode float8 to float32.""" 

98 if encoded == 0: 

99 return 0.0 

100 

101 # Extract components 

102 sign = (encoded >> 7) & 1 

103 exp = (encoded >> self.mantissa_bits) & self.exponent_mask 

104 mant = encoded & self.mantissa_mask 

105 

106 # Handle special cases 

107 if exp == self.exponent_mask: 

108 if mant == 0: 

109 return -np.inf if sign else np.inf 

110 else: 

111 # NaN 

112 return np.nan 

113 

114 # Denormalized numbers 

115 if exp == 0: 

116 if mant == 0: 

117 return -0.0 if sign else 0.0 

118 else: 

119 # Denormalized: (-1)^sign * 2^(-bias+1) * mant/2^mantissa_bits 

120 value = (mant / (1 << self.mantissa_bits)) * (2 ** (-self.bias + 1)) 

121 return -value if sign else value 

122 

123 # Normalized numbers: (-1)^sign * 2^(exp-bias) * (1 + mant/2^mantissa_bits) 

124 value = (1 + mant / (1 << self.mantissa_bits)) * (2 ** (exp - self.bias)) 

125 return -value if sign else value 

126 

127 def encode_array(self, arr: np.ndarray) -> np.ndarray: 

128 """Encode numpy array to float8.""" 

129 encoded = np.zeros(arr.shape, dtype=np.uint8) 

130 flat_arr = arr.flatten() 

131 flat_encoded = encoded.flatten() 

132 

133 for i in range(flat_arr.size): 

134 flat_encoded[i] = self.encode_scalar(float(flat_arr[i])) 

135 

136 return flat_encoded.reshape(arr.shape) 

137 

138 def decode_array(self, encoded: np.ndarray) -> np.ndarray: 

139 """Decode float8 array to float32.""" 

140 decoded = np.zeros(encoded.shape, dtype=np.float32) 

141 flat_encoded = encoded.flatten() 

142 flat_decoded = decoded.flatten() 

143 

144 for i in range(flat_encoded.size): 

145 flat_decoded[i] = self.decode_scalar(int(flat_encoded[i])) 

146 

147 return decoded.reshape(encoded.shape) 

148 

149 

150def pack_float8_array(arr: np.ndarray, format_type: str = "e4m3") -> tuple[np.ndarray, Float8Encoder]: 

151 """ 

152 Pack float32 array into float8 format. 

153  

154 Args: 

155 arr: Input float32 array 

156 format_type: 'e4m3' or 'e5m2' 

157  

158 Returns: 

159 Tuple of (packed_uint8_array, encoder) 

160 """ 

161 encoder = Float8Encoder(format_type) 

162 return encoder.encode_array(arr.astype(np.float32)), encoder 

163 

164 

165def unpack_float8_array(packed: np.ndarray, encoder: Float8Encoder) -> np.ndarray: 

166 """ 

167 Unpack float8 array to float32. 

168  

169 Args: 

170 packed: Packed uint8 array 

171 encoder: Float8Encoder instance 

172  

173 Returns: 

174 Unpacked float32 array 

175 """ 

176 return encoder.decode_array(packed)