Skip to content

Commit c9351d6

Browse files
Use qd.interop.get_mps_command_queue() instead of inline extraction
The ctypes/ObjC logic to extract PyTorch MPS's MTLCommandQueue* now lives in Quadrants (qd.interop module), so Genesis just calls the public API. Removes ~35 lines of low-level ctypes code. Co-authored-by: Cursor <cursoragent@cursor.com>
1 parent 3d0330c commit c9351d6

1 file changed

Lines changed: 1 addition & 37 deletions

File tree

genesis/__init__.py

Lines changed: 1 addition & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
import ctypes
21
import io
32
import os
43
import sys
@@ -53,41 +52,6 @@
5352
EPS: float | None = None
5453

5554

56-
def _get_mps_command_queue() -> int:
57-
"""Extract PyTorch MPS's MTLCommandQueue* as a Python int, or 0 on failure."""
58-
try:
59-
torch_lib = os.path.join(os.path.dirname(torch.__file__), "lib", "libtorch_cpu.dylib")
60-
handle = ctypes.CDLL(torch_lib)._handle
61-
62-
libdl = ctypes.CDLL(None)
63-
dlsym = libdl.dlsym
64-
dlsym.restype = ctypes.c_void_p
65-
dlsym.argtypes = [ctypes.c_void_p, ctypes.c_char_p]
66-
67-
stream_fn = dlsym(handle, b"_ZN2at3mps19getDefaultMPSStreamEv")
68-
if not stream_fn:
69-
return 0
70-
stream_ptr = ctypes.CFUNCTYPE(ctypes.c_void_p)(stream_fn)()
71-
72-
cb_fn = dlsym(handle, b"_ZN2at3mps9MPSStream13commandBufferEv")
73-
if not cb_fn:
74-
return 0
75-
cb_ptr = ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.c_void_p)(cb_fn)(stream_ptr)
76-
77-
objc = ctypes.CDLL("/usr/lib/libobjc.A.dylib")
78-
sel_reg = objc.sel_registerName
79-
sel_reg.restype = ctypes.c_void_p
80-
sel_reg.argtypes = [ctypes.c_char_p]
81-
msg_send = objc.objc_msgSend
82-
msg_send.restype = ctypes.c_void_p
83-
msg_send.argtypes = [ctypes.c_void_p, ctypes.c_void_p]
84-
85-
queue_ptr = msg_send(cb_ptr, sel_reg(b"commandQueue"))
86-
return queue_ptr or 0
87-
except (OSError, AttributeError):
88-
return 0
89-
90-
9155
########################## init ##########################
9256
def init(
9357
*,
@@ -294,7 +258,7 @@ def init(
294258
# On Metal, share PyTorch MPS's command queue with Quadrants so that GPU-side ordering is automatic and the
295259
# per-interop-point sync overhead (qd.sync / torch.mps.synchronize) is eliminated.
296260
if backend == _gs_backend.metal and device.type == "mps":
297-
mps_queue = _get_mps_command_queue()
261+
mps_queue = qd.interop.get_mps_command_queue()
298262
if not mps_queue:
299263
raise_exception(
300264
"Failed to extract PyTorch MPS's Metal command queue. "

0 commit comments

Comments
 (0)