Coverage for src/pymetallic/metallic.py: 83%

1import ctypes 

2import os 

3from ctypes import c_void_p, c_char_p, c_int32, c_uint64, c_bool 

4from typing import List, Optional, Tuple 

5 

6import numpy as np 

7from numpy.typing import DTypeLike 

8 

9 

10# Public error type 

11class MetalError(Exception): 

12 pass 

13 

14 

15# Internal FFI loader 

16_metal_lib = None 

17 

18 

19def _load_metal_library() -> ctypes.CDLL: 

20 # Prefer the packaged dylib sitting next to this file 

21 pkg_dir = os.path.dirname(os.path.abspath(__file__)) 

22 local_path = os.path.join(pkg_dir, "libpymetallic.dylib") 

23 if os.path.exists(local_path): 

24 return ctypes.CDLL(local_path) 

25 # Fallback to default loader if needed 

26 try: 

27 return ctypes.CDLL("libpymetallic.dylib") 

28 except OSError as e: 

29 raise MetalError( 

30 "Failed to load libpymetallic.dylib. Ensure it is built and on your DYLD_LIBRARY_PATH" 

31 ) from e 

32 

33 

34def _setup_function_signatures(lib: ctypes.CDLL) -> None: 

35 # Devices 

36 lib.metal_get_default_device.restype = c_void_p 

37 lib.metal_get_all_devices.restype = ctypes.POINTER(c_void_p) 

38 lib.metal_get_device_count.restype = ctypes.c_int32 

39 lib.metal_device_get_name.argtypes = [c_void_p] 

40 lib.metal_device_get_name.restype = c_char_p 

41 lib.metal_device_supports_shader_barycentric_coordinates.argtypes = [c_void_p] 

42 lib.metal_device_supports_shader_barycentric_coordinates.restype = c_bool 

43 

44 # Command queue and buffer 

45 lib.metal_device_make_command_queue.argtypes = [c_void_p] 

46 lib.metal_device_make_command_queue.restype = c_void_p 

47 lib.metal_command_queue_make_command_buffer.argtypes = [c_void_p] 

48 lib.metal_command_queue_make_command_buffer.restype = c_void_p 

49 

50 # Buffers 

51 lib.metal_device_make_buffer.argtypes = [c_void_p, c_uint64, c_int32] 

52 lib.metal_device_make_buffer.restype = c_void_p 

53 lib.metal_device_make_buffer_with_bytes.argtypes = [ 

54 c_void_p, 

55 c_void_p, 

56 c_uint64, 

57 c_int32, 

58 ] 

59 lib.metal_device_make_buffer_with_bytes.restype = c_void_p 

60 lib.metal_buffer_get_contents.argtypes = [c_void_p] 

61 lib.metal_buffer_get_contents.restype = c_void_p 

62 lib.metal_buffer_get_length.argtypes = [c_void_p] 

63 lib.metal_buffer_get_length.restype = c_uint64 

64 

65 # Library/functions 

66 lib.metal_device_make_library_with_source.argtypes = [c_void_p, c_char_p] 

67 lib.metal_device_make_library_with_source.restype = c_void_p 

68 lib.metal_library_make_function.argtypes = [c_void_p, c_char_p] 

69 lib.metal_library_make_function.restype = c_void_p 

70 

71 # Compute pipeline 

72 lib.metal_device_make_compute_pipeline_state.argtypes = [c_void_p, c_void_p] 

73 lib.metal_device_make_compute_pipeline_state.restype = c_void_p 

74 

75 # Encoders 

76 lib.metal_command_buffer_make_compute_command_encoder.argtypes = [c_void_p] 

77 lib.metal_command_buffer_make_compute_command_encoder.restype = c_void_p 

78 lib.metal_compute_command_encoder_set_compute_pipeline_state.argtypes = [ 

79 c_void_p, 

80 c_void_p, 

81 ] 

82 lib.metal_compute_command_encoder_set_compute_pipeline_state.restype = None 

83 lib.metal_compute_command_encoder_set_buffer.argtypes = [ 

84 c_void_p, 

85 c_void_p, 

86 c_uint64, 

87 c_int32, 

88 ] 

89 lib.metal_compute_command_encoder_set_buffer.restype = None 

90 lib.metal_compute_command_encoder_set_threadgroup_memory_length.argtypes = [ 

91 c_void_p, 

92 c_uint64, 

93 c_int32, 

94 ] 

95 lib.metal_compute_command_encoder_set_threadgroup_memory_length.restype = None 

96 lib.metal_compute_command_encoder_dispatch_threads.argtypes = [ 

97 c_void_p, 

98 c_int32, 

99 c_int32, 

100 c_int32, 

101 c_int32, 

102 c_int32, 

103 c_int32, 

104 ] 

105 lib.metal_compute_command_encoder_dispatch_threads.restype = None 

106 lib.metal_compute_command_encoder_dispatch_threadgroups.argtypes = [ 

107 c_void_p, 

108 c_int32, 

109 c_int32, 

110 c_int32, 

111 c_int32, 

112 c_int32, 

113 c_int32, 

114 ] 

115 lib.metal_compute_command_encoder_dispatch_threadgroups.restype = None 

116 lib.metal_compute_command_encoder_end_encoding.argtypes = [c_void_p] 

117 lib.metal_compute_command_encoder_end_encoding.restype = None 

118 

119 # Command buffer exec 

120 lib.metal_command_buffer_commit.argtypes = [c_void_p] 

121 lib.metal_command_buffer_commit.restype = None 

122 lib.metal_command_buffer_wait_until_completed.argtypes = [c_void_p] 

123 lib.metal_command_buffer_wait_until_completed.restype = None 

124 

125 

126def _get_metal_lib() -> ctypes.CDLL: 

127 global _metal_lib 

128 if _metal_lib is None: 

129 _metal_lib = _load_metal_library() 

130 _setup_function_signatures(_metal_lib) 

131 return _metal_lib 

132 

133 

134# Public storage mode constants 

135 

136 

137class Device: 

138 """Metal device wrapper""" 

139 

140 STORAGE_SHARED = 0 

141 

142 def __init__(self, device_ptr: int): 

143 self._device_ptr = c_void_p(device_ptr) 

144 

145 @classmethod 

146 def get_default_device(cls) -> "Device": 

147 lib = _get_metal_lib() 

148 device_ptr = lib.metal_get_default_device() 

149 if not device_ptr: 

150 raise MetalError("No Metal devices available") 

151 return cls(device_ptr) 

152 

153 @classmethod 

154 def get_all_devices(cls) -> List["Device"]: 

155 lib = _get_metal_lib() 

156 count = lib.metal_get_device_count() 

157 if count <= 0: 

158 return [] 

159 devices_ptr = lib.metal_get_all_devices() 

160 result: List[Device] = [] 

161 # devices_ptr is a pointer to an array of c_void_p of length count 

162 for i in range(count): 

163 ptr = devices_ptr[i] 

164 if ptr: 

165 result.append(cls(ptr)) 

166 return result 

167 

168 @property 

169 def name(self) -> str: 

170 lib = _get_metal_lib() 

171 name_ptr = lib.metal_device_get_name(self._device_ptr) 

172 return name_ptr.decode("utf-8") if name_ptr else "Unknown Device" 

173 

174 def supports_shader_barycentric_coordinates(self) -> bool: 

175 lib = _get_metal_lib() 

176 return bool( 

177 lib.metal_device_supports_shader_barycentric_coordinates(self._device_ptr) 

178 ) 

179 

180 def make_compute_pipeline_state(self, fun: "Function") -> "ComputePipelineState": 

181 return ComputePipelineState(self, fun) 

182 

183 def make_buffer_from_numpy(self, array, storage_mode: int | None = None): 

184 storage_mode = Buffer.STORAGE_SHARED if storage_mode is None else storage_mode 

185 return Buffer.from_numpy(self, array, storage_mode) 

186 

187 def make_buffer(self, size: int, storage_mode: int | None = None): 

188 storage_mode = Buffer.STORAGE_SHARED if storage_mode is None else storage_mode 

189 return Buffer(self, size, storage_mode) 

190 

191 def make_command_queue(self): 

192 return CommandQueue(self) 

193 

194 def make_library(self, source): 

195 return Library(self, source) 

196 

197 

198class CommandQueue: 

199 """Metal command queue wrapper""" 

200 

201 def __init__(self, device: Device): 

202 self.device = device 

203 lib = _get_metal_lib() 

204 self._queue_ptr = c_void_p( 

205 lib.metal_device_make_command_queue(device._device_ptr) 

206 ) 

207 if not self._queue_ptr: 

208 raise MetalError("Failed to create command queue") 

209 

210 def make_command_buffer(self) -> "CommandBuffer": 

211 return CommandBuffer(self) 

212 

213 

214class Buffer: 

215 """Metal buffer wrapper - similar to pyopencl.Buffer""" 

216 

217 STORAGE_SHARED = 0 

218 STORAGE_MANAGED = 1 

219 STORAGE_PRIVATE = 2 

220 

221 @staticmethod 

222 def _storage_mode_to_resource_options(storage_mode: int) -> int: 

223 # Map to MTLResourceOptions bitfield; default CPU cache (0), storage mode shifted by 4 

224 STORAGE_SHIFT = 4 

225 if storage_mode == Buffer.STORAGE_SHARED: 

226 return 0 

227 if storage_mode == Buffer.STORAGE_MANAGED: 

228 return 1 << STORAGE_SHIFT 

229 if storage_mode == Buffer.STORAGE_PRIVATE: 

230 return 2 << STORAGE_SHIFT 

231 raise ValueError( 

232 f"Invalid storage_mode {storage_mode}. " 

233 f"Use Buffer.STORAGE_SHARED, STORAGE_MANAGED, or STORAGE_PRIVATE." 

234 ) 

235 

236 def __init__(self, device: Device, size: int, storage_mode: int = STORAGE_SHARED): 

237 self.device = device 

238 self.size = size 

239 lib = _get_metal_lib() 

240 resource_options = self._storage_mode_to_resource_options(storage_mode) 

241 self._buffer_ptr = c_void_p( 

242 lib.metal_device_make_buffer( 

243 device._device_ptr, c_uint64(size), c_int32(resource_options) 

244 ) 

245 ) 

246 if not self._buffer_ptr: 

247 raise MetalError("Failed to create buffer") 

248 

249 @classmethod 

250 def from_numpy( 

251 cls, device: Device, array: np.ndarray, storage_mode: int = STORAGE_SHARED 

252 ) -> "Buffer": 

253 lib = _get_metal_lib() 

254 resource_options = cls._storage_mode_to_resource_options(storage_mode) 

255 arr = np.ascontiguousarray(array) 

256 ptr = lib.metal_device_make_buffer_with_bytes( 

257 device._device_ptr, 

258 arr.ctypes.data_as(c_void_p), 

259 c_uint64(arr.nbytes), 

260 c_int32(resource_options), 

261 ) 

262 if not ptr: 

263 raise MetalError("Failed to create buffer from numpy array") 

264 buf = cls.__new__(cls) 

265 buf.device = device 

266 buf.size = arr.nbytes 

267 buf._buffer_ptr = c_void_p(ptr) 

268 return buf 

269 

270 def to_numpy( 

271 self, dtype: DTypeLike, shape: Optional[Tuple[int, ...]] = None 

272 ) -> np.ndarray: 

273 lib = _get_metal_lib() 

274 contents_ptr = lib.metal_buffer_get_contents(self._buffer_ptr) 

275 if not contents_ptr: 

276 raise MetalError( 

277 "Buffer has no CPU-accessible contents (likely private storage)" 

278 ) 

279 length = lib.metal_buffer_get_length(self._buffer_ptr) 

280 

281 # Correctly obtain the integer address of the buffer's contents 

282 addr = ctypes.cast(contents_ptr, c_void_p).value 

283 if addr is None: 

284 raise MetalError("Failed to obtain buffer contents address") 

285 

286 # Build a ctypes array that views the raw bytes at the given address 

287 length_int = int(length) 

288 raw = (ctypes.c_ubyte * length_int).from_address(addr) 

289 

290 # Create a NumPy view over the raw bytes 

291 buffer_array = np.frombuffer(raw, dtype=np.uint8, count=length_int) 

292 

293 # Validate dtype alignment 

294 target_dtype = np.dtype(dtype) 

295 if length_int % target_dtype.itemsize != 0: 

296 raise MetalError( 

297 f"Buffer length {length_int} is not a multiple of dtype itemsize {target_dtype.itemsize}" 

298 ) 

299 

300 typed = buffer_array.view(dtype=target_dtype) 

301 

302 # Validate and apply shape if provided 

303 if shape is not None: 

304 expected_elems = 1 

305 for s in shape: 

306 expected_elems *= int(s) 

307 actual_elems = length_int // target_dtype.itemsize 

308 if expected_elems != actual_elems: 

309 raise MetalError( 

310 f"Shape {shape} has {expected_elems} elements, but buffer holds {actual_elems}" 

311 ) 

312 return typed.reshape(shape) 

313 

314 return typed 

315 

316 

317class Library: 

318 """Metal shader library wrapper""" 

319 

320 def __init__(self, device: Device, source: str): 

321 self.device = device 

322 lib = _get_metal_lib() 

323 self._library_ptr = c_void_p( 

324 lib.metal_device_make_library_with_source( 

325 device._device_ptr, source.encode("utf-8") 

326 ) 

327 ) 

328 if not self._library_ptr: 

329 raise MetalError("Failed to compile Metal library") 

330 

331 def make_function(self, name: str) -> "Function": 

332 lib = _get_metal_lib() 

333 func_ptr = lib.metal_library_make_function( 

334 self._library_ptr, name.encode("utf-8") 

335 ) 

336 if not func_ptr: 

337 raise MetalError(f"Function '{name}' not found in library") 

338 return Function(self, name, func_ptr) 

339 

340 

341class Function: 

342 """Compiled Metal function wrapper""" 

343 

344 def __init__(self, library: Library, name: str, function_ptr: int): 

345 self.library = library 

346 self.name = name 

347 self._function_ptr = c_void_p(function_ptr) 

348 

349 

350class ComputePipelineState: 

351 """Compute pipeline state wrapper""" 

352 

353 def __init__(self, device: Device, function: Function): 

354 self.device = device 

355 self.function = function 

356 lib = _get_metal_lib() 

357 self._pipeline_ptr = c_void_p( 

358 lib.metal_device_make_compute_pipeline_state( 

359 device._device_ptr, function._function_ptr 

360 ) 

361 ) 

362 if not self._pipeline_ptr: 

363 raise MetalError("Failed to create compute pipeline state") 

364 

365 

366class CommandBuffer: 

367 """Metal command buffer wrapper""" 

368 

369 def __init__(self, command_queue: CommandQueue): 

370 self.command_queue = command_queue 

371 lib = _get_metal_lib() 

372 self._buffer_ptr = c_void_p( 

373 lib.metal_command_queue_make_command_buffer(command_queue._queue_ptr) 

374 ) 

375 if not self._buffer_ptr: 

376 raise MetalError("Failed to create command buffer") 

377 

378 def make_compute_command_encoder(self) -> "ComputeCommandEncoder": 

379 return ComputeCommandEncoder(self) 

380 

381 def commit(self): 

382 lib = _get_metal_lib() 

383 lib.metal_command_buffer_commit(self._buffer_ptr) 

384 

385 def wait_until_completed(self): 

386 lib = _get_metal_lib() 

387 lib.metal_command_buffer_wait_until_completed(self._buffer_ptr) 

388 

389 def compute_command_encoder(self): 

390 return ComputeCommandEncoder(self) 

391 

392 

393class ComputeCommandEncoder: 

394 """Metal compute command encoder""" 

395 

396 def __init__(self, command_buffer: CommandBuffer): 

397 self.command_buffer = command_buffer 

398 lib = _get_metal_lib() 

399 self._encoder_ptr = c_void_p( 

400 lib.metal_command_buffer_make_compute_command_encoder( 

401 command_buffer._buffer_ptr 

402 ) 

403 ) 

404 if not self._encoder_ptr: 

405 raise MetalError("Failed to create compute command encoder") 

406 

407 def set_compute_pipeline_state(self, pipeline: ComputePipelineState): 

408 lib = _get_metal_lib() 

409 lib.metal_compute_command_encoder_set_compute_pipeline_state( 

410 self._encoder_ptr, pipeline._pipeline_ptr 

411 ) 

412 

413 def set_buffer(self, buffer: Buffer, offset: int, index: int): 

414 lib = _get_metal_lib() 

415 lib.metal_compute_command_encoder_set_buffer( 

416 self._encoder_ptr, buffer._buffer_ptr, c_uint64(offset), c_int32(index) 

417 ) 

418 

419 def set_threadgroup_memory_length(self, length: int, index: int): 

420 lib = _get_metal_lib() 

421 lib.metal_compute_command_encoder_set_threadgroup_memory_length( 

422 self._encoder_ptr, c_uint64(length), c_int32(index) 

423 ) 

424 

425 def dispatch_threads( 

426 self, 

427 threads_per_grid: Tuple[int, int, int], 

428 threads_per_threadgroup: Tuple[int, int, int], 

429 ): 

430 lib = _get_metal_lib() 

431 tx, ty, tz = map(int, threads_per_grid) 

432 tgx, tgy, tgz = map(int, threads_per_threadgroup) 

433 lib.metal_compute_command_encoder_dispatch_threads( 

434 self._encoder_ptr, tx, ty, tz, tgx, tgy, tgz 

435 ) 

436 

437 def dispatch_threadgroups( 

438 self, 

439 groups: Tuple[int, int, int], 

440 threads_per_threadgroup: Tuple[int, int, int], 

441 ): 

442 lib = _get_metal_lib() 

443 gx, gy, gz = map(int, groups) 

444 tx, ty, tz = map(int, threads_per_threadgroup) 

445 lib.metal_compute_command_encoder_dispatch_threadgroups( 

446 self._encoder_ptr, gx, gy, gz, tx, ty, tz 

447 ) 

448 

449 def end_encoding(self): 

450 lib = _get_metal_lib() 

451 lib.metal_compute_command_encoder_end_encoding(self._encoder_ptr) 

452 

453 

454def run_simple_compute_example(): 

455 """Small self-test: vector add""" 

456 device = Device.get_default_device() 

457 queue = device.make_command_queue() 

458 

459 n = 256 

460 a = np.random.rand(n).astype(np.float32) 

461 b = np.random.rand(n).astype(np.float32) 

462 

463 buf_a = device.make_buffer_from_numpy(a) 

464 buf_b = device.make_buffer_from_numpy(b) 

465 buf_out = device.make_buffer(n * 4) 

466 

467 source = """ 

468 #include <metal_stdlib> 

469 using namespace metal; 

470 kernel void vector_add(device const float* a [[buffer(0)]], 

471 device const float* b [[buffer(1)]], 

472 device float* out [[buffer(2)]], 

473 uint idx [[thread_position_in_grid]]) { 

474 out[idx] = a[idx] + b[idx]; 

475 }""" 

476 

477 lib = device.make_library(source) 

478 fn = lib.make_function("vector_add") 

479 pso = ComputePipelineState(device, fn) 

480 

481 cb = queue.make_command_buffer() 

482 enc = cb.make_compute_command_encoder() 

483 enc.set_compute_pipeline_state(pso) 

484 enc.set_buffer(buf_a, 0, 0) 

485 enc.set_buffer(buf_b, 0, 1) 

486 enc.set_buffer(buf_out, 0, 2) 

487 enc.dispatch_threads((n, 1, 1), (64, 1, 1)) 

488 enc.end_encoding() 

489 cb.commit() 

490 cb.wait_until_completed() 

491 

492 out = buf_out.to_numpy(np.float32, (n,)) 

493 assert np.allclose(out, a + b, rtol=1e-5), "Vector add result mismatch" 

494 print("Simple compute example succeeded on:", device.name)