Coverage for src / monte_neo / core / native / metal_engine.py: 13%

1""" 

2Metal GPU acceleration engine for Apple Silicon. 

3 

4Provides direct Metal shader access for maximum performance. 

5""" 

6 

7from __future__ import annotations 

8 

9import os 

10from typing import TYPE_CHECKING 

11 

12import numpy as np 

13 

14if TYPE_CHECKING: 

15 pass 

16 

17# Try to import Metal framework 

18try: 

19 import Metal 

20 METAL_AVAILABLE = True 

21except ImportError: 

22 METAL_AVAILABLE = False 

23 

24 

25class MetalFloat8Engine: 

26 """Metal-based GPU engine for float8 operations.""" 

27 

28 def __init__(self) -> None: 

29 """Initialize Metal engine.""" 

30 if not METAL_AVAILABLE: 

31 raise RuntimeError("Metal framework not available. Install pyobjc-framework-Metal.") 

32 

33 self.device = Metal.MTLCreateSystemDefaultDevice() 

34 if self.device is None: 

35 raise RuntimeError("No Metal device available.") 

36 

37 self.command_queue = self.device.newCommandQueue() 

38 

39 # Load shaders 

40 self._load_shaders() 

41 

42 def _load_shaders(self) -> None: 

43 """Load Metal shaders from file.""" 

44 shader_path = os.path.join(os.path.dirname(__file__), "metal_kernels.metal") 

45 

46 if not os.path.exists(shader_path): 

47 # Create default shaders if file doesn't exist 

48 self._create_default_shaders() 

49 return 

50 

51 with open(shader_path) as f: 

52 shader_source = f.read() 

53 

54 # Compile shader 

55 library, error = self.device.newLibraryWithSource_options_error_( 

56 shader_source, None, None 

57 ) 

58 

59 if library is None: 

60 raise RuntimeError(f"Failed to compile Metal shaders: {error}") 

61 

62 self.library = library 

63 

64 # Get function handles 

65 self.float32_to_e4m3_func = library.newFunctionWithName_("float32_to_float8_e4m3") 

66 self.e4m3_to_float32_func = library.newFunctionWithName_("float8_e4m3_to_float32") 

67 self.float32_to_e5m2_func = library.newFunctionWithName_("float32_to_float8_e5m2") 

68 self.e5m2_to_float32_func = library.newFunctionWithName_("float8_e5m2_to_float32") 

69 self.generate_scenarios_func = library.newFunctionWithName_("generate_shuffle_scenarios_coalesced") 

70 

71 def _create_default_shaders(self) -> None: 

72 """Create default compute pipeline for basic operations.""" 

73 # This is a fallback - in production, use pre-compiled shaders 

74 shader_source = """ 

75 #include <metal_stdlib> 

76 using namespace metal; 

77  

78 kernel void float32_to_float8_e4m3( 

79 const device float* input [[buffer(0)]], 

80 device uchar* output [[buffer(1)]], 

81 uint id [[thread_position_in_grid]]) 

82 { 

83 // Simple quantization for testing 

84 float val = input[id]; 

85 uchar quantized = uchar(clamp(val * 127.0, -128.0, 127.0) + 128.0); 

86 output[id] = quantized; 

87 } 

88  

89 kernel void float8_e4m3_to_float32( 

90 const device uchar* input [[buffer(0)]], 

91 device float* output [[buffer(1)]], 

92 uint id [[thread_position_in_grid]]) 

93 { 

94 uchar val = input[id]; 

95 output[id] = (float(val) - 128.0) / 127.0; 

96 } 

97 """ 

98 

99 library, error = self.device.newLibraryWithSource_options_error_( 

100 shader_source, None, None 

101 ) 

102 

103 if library is None: 

104 raise RuntimeError(f"Failed to create default Metal shaders: {error}") 

105 

106 self.library = library 

107 self.float32_to_e4m3_func = library.newFunctionWithName_("float32_to_float8_e4m3") 

108 self.e4m3_to_float32_func = library.newFunctionWithName_("float8_e4m3_to_float32") 

109 

110 def encode_float32_to_e4m3(self, input_array: np.ndarray) -> np.ndarray: 

111 """Encode float32 array to float8 E4M3 format.""" 

112 if self.float32_to_e4m3_func is None: 

113 raise RuntimeError("E4M3 conversion function not available") 

114 

115 # Flatten array for processing 

116 flat_input = input_array.flatten().astype(np.float32) 

117 n_elements = len(flat_input) 

118 

119 # Create output array 

120 output_array = np.zeros(n_elements, dtype=np.uint8) 

121 

122 # Create Metal buffers 

123 input_buffer = self.device.newBufferWithBytes_length_options_( 

124 flat_input, n_elements * 4, Metal.MTLResourceStorageModeShared 

125 ) 

126 output_buffer = self.device.newBufferWithLength_options_(n_elements, Metal.MTLResourceStorageModeShared) 

127 

128 # Create compute pipeline 

129 pipeline, error = self.device.newComputePipelineStateWithFunction_error_( 

130 self.float32_to_e4m3_func, None 

131 ) 

132 

133 if pipeline is None: 

134 raise RuntimeError(f"Failed to create pipeline: {error}") 

135 

136 # Create command encoder 

137 command_buffer = self.command_queue.commandBuffer() 

138 encoder = command_buffer.computeCommandEncoder() 

139 encoder.setComputePipelineState_(pipeline) 

140 

141 # Set buffers 

142 encoder.setBuffer_offset_atIndex_(input_buffer, 0, 0) 

143 encoder.setBuffer_offset_atIndex_(output_buffer, 0, 1) 

144 

145 # Dispatch threads 

146 threads_per_threadgroup = pipeline.maxTotalThreadsPerThreadgroup() 

147 threadgroups = (n_elements + threads_per_threadgroup - 1) // threads_per_threadgroup 

148 

149 encoder.dispatchThreadgroups_threadsPerThreadgroup_( 

150 (threadgroups, 1, 1), (threads_per_threadgroup, 1, 1) 

151 ) 

152 

153 encoder.endEncoding() 

154 command_buffer.commit() 

155 command_buffer.waitUntilCompleted() 

156 

157 # Copy result back 

158 output_array[:] = np.frombuffer( 

159 output_buffer.contents().as_buffer(n_elements), 

160 dtype=np.uint8 

161 ) 

162 

163 return output_array.reshape(input_array.shape) 

164 

165 def decode_e4m3_to_float32(self, input_array: np.ndarray) -> np.ndarray: 

166 """Decode float8 E4M3 array to float32.""" 

167 if self.e4m3_to_float32_func is None: 

168 raise RuntimeError("E4M3 decoding function not available") 

169 

170 # Flatten array for processing 

171 flat_input = input_array.flatten().astype(np.uint8) 

172 n_elements = len(flat_input) 

173 

174 # Create output array 

175 output_array = np.zeros(n_elements, dtype=np.float32) 

176 

177 # Create Metal buffers 

178 input_buffer = self.device.newBufferWithBytes_length_options_( 

179 flat_input, n_elements, Metal.MTLResourceStorageModeShared 

180 ) 

181 output_buffer = self.device.newBufferWithLength_options_(n_elements * 4, Metal.MTLResourceStorageModeShared) 

182 

183 # Create compute pipeline 

184 pipeline, error = self.device.newComputePipelineStateWithFunction_error_( 

185 self.e4m3_to_float32_func, None 

186 ) 

187 

188 if pipeline is None: 

189 raise RuntimeError(f"Failed to create pipeline: {error}") 

190 

191 # Create command encoder 

192 command_buffer = self.command_queue.commandBuffer() 

193 encoder = command_buffer.computeCommandEncoder() 

194 encoder.setComputePipelineState_(pipeline) 

195 

196 # Set buffers 

197 encoder.setBuffer_offset_atIndex_(input_buffer, 0, 0) 

198 encoder.setBuffer_offset_atIndex_(output_buffer, 0, 1) 

199 

200 # Dispatch threads 

201 threads_per_threadgroup = pipeline.maxTotalThreadsPerThreadgroup() 

202 threadgroups = (n_elements + threads_per_threadgroup - 1) // threads_per_threadgroup 

203 

204 encoder.dispatchThreadgroups_threadsPerThreadgroup_( 

205 (threadgroups, 1, 1), (threads_per_threadgroup, 1, 1) 

206 ) 

207 

208 encoder.endEncoding() 

209 command_buffer.commit() 

210 command_buffer.waitUntilCompleted() 

211 

212 # Copy result back 

213 output_array[:] = np.frombuffer( 

214 output_buffer.contents().as_buffer(n_elements * 4), 

215 dtype=np.float32 

216 ) 

217 

218 return output_array.reshape(input_array.shape) 

219 

220 def generate_scenarios_e4m3(self, base_prices: np.ndarray, n_scenarios: int, seed: int = 42) -> np.ndarray: 

221 """Generate scenarios using Metal.""" 

222 if self.generate_scenarios_func is None: 

223 raise RuntimeError("Scenario generation function not available") 

224 

225 time_steps = len(base_prices) 

226 n_elements = n_scenarios * time_steps 

227 

228 # Generate random indices (N x T-1) 

229 np.random.seed(seed) 

230 random_indices = np.random.randint(0, time_steps - 1, (n_scenarios, time_steps - 1)).astype(np.uint32) 

231 

232 # Create output array 

233 output_array = np.zeros(n_elements, dtype=np.uint8) 

234 

235 # Create Metal buffers 

236 base_prices_buffer = self.device.newBufferWithBytes_length_options_( 

237 base_prices, time_steps, Metal.MTLResourceStorageModeShared 

238 ) 

239 scenarios_buffer = self.device.newBufferWithLength_options_(n_elements, Metal.MTLResourceStorageModeShared) 

240 indices_buffer = self.device.newBufferWithBytes_length_options_( 

241 random_indices, random_indices.nbytes, Metal.MTLResourceStorageModeShared 

242 ) 

243 

244 # Create compute pipeline 

245 pipeline, error = self.device.newComputePipelineStateWithFunction_error_( 

246 self.generate_scenarios_func, None 

247 ) 

248 

249 if pipeline is None: 

250 raise RuntimeError(f"Failed to create pipeline: {error}") 

251 

252 # Create command encoder 

253 command_buffer = self.command_queue.commandBuffer() 

254 encoder = command_buffer.computeCommandEncoder() 

255 encoder.setComputePipelineState_(pipeline) 

256 

257 # Set buffers 

258 encoder.setBuffer_offset_atIndex_(base_prices_buffer, 0, 0) 

259 encoder.setBuffer_offset_atIndex_(scenarios_buffer, 0, 1) 

260 encoder.setBuffer_offset_atIndex_(indices_buffer, 0, 2) 

261 

262 # Set constants 

263 # Using bytes to pass uint32 values to Metal 

264 n_scen_bytes = n_scenarios.to_bytes(4, byteorder='little') 

265 t_steps_bytes = time_steps.to_bytes(4, byteorder='little') 

266 

267 encoder.setBytes_length_atIndex_(n_scen_bytes, 4, 3) 

268 encoder.setBytes_length_atIndex_(t_steps_bytes, 4, 4) 

269 

270 # Dispatch threads 

271 threads_per_threadgroup = pipeline.maxTotalThreadsPerThreadgroup() 

272 threadgroups = (n_elements + threads_per_threadgroup - 1) // threads_per_threadgroup 

273 

274 encoder.dispatchThreadgroups_threadsPerThreadgroup_( 

275 (threadgroups, 1, 1), (threads_per_threadgroup, 1, 1) 

276 ) 

277 

278 encoder.endEncoding() 

279 command_buffer.commit() 

280 command_buffer.waitUntilCompleted() 

281 

282 # Copy result back 

283 output_array[:] = np.frombuffer( 

284 scenarios_buffer.contents().as_buffer(n_elements), 

285 dtype=np.uint8 

286 ) 

287 

288 return output_array.reshape((n_scenarios, time_steps)) 

289 

290 def get_memory_bandwidth_improvement(self) -> float: 

291 return 4.0 

292 

293 def get_precision_loss_estimate(self) -> float: 

294 return 4.0