|
1 | | -import ctypes |
2 | 1 | import io |
3 | 2 | import os |
4 | 3 | import sys |
|
53 | 52 | EPS: float | None = None |
54 | 53 |
|
55 | 54 |
|
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 | | - |
91 | 55 | ########################## init ########################## |
92 | 56 | def init( |
93 | 57 | *, |
@@ -294,7 +258,7 @@ def init( |
294 | 258 | # On Metal, share PyTorch MPS's command queue with Quadrants so that GPU-side ordering is automatic and the |
295 | 259 | # per-interop-point sync overhead (qd.sync / torch.mps.synchronize) is eliminated. |
296 | 260 | 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() |
298 | 262 | if not mps_queue: |
299 | 263 | raise_exception( |
300 | 264 | "Failed to extract PyTorch MPS's Metal command queue. " |
|
0 commit comments