[FEATURE] Share Metal command queue between Quadrants and PyTorch MPS

At gs.init() time on Apple Metal, extract PyTorch MPS's MTLCommandQueue
and pass it to Quadrants via external_metal_command_queue. This lets both
frameworks dispatch GPU work on the same queue, so Metal's sequential
command buffer ordering eliminates the need for explicit CPU-side syncs
(qd.sync / torch.mps.synchronize) at every interop point.

Raises an exception if the queue cannot be extracted, since correct
synchronisation on Metal depends on the shared queue.

Removes all torch.mps.synchronize() guards from solvers and misc.py.

Depends on Genesis-Embodied-AI/quadrants#618.

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: hughperkins

genesis/__init__.py

@@ -1,3 +1,4 @@
+import ctypes
 import io
 import os
 import sys
@@ -52,6 +53,41 @@
 EPS: float | None = None
 
 
+def _get_mps_command_queue() -> int:
+    """Extract PyTorch MPS's MTLCommandQueue* as a Python int, or 0 on failure."""
+    try:
+        torch_lib = os.path.join(os.path.dirname(torch.__file__), "lib", "libtorch_cpu.dylib")
+        handle = ctypes.CDLL(torch_lib)._handle
+
+        libdl = ctypes.CDLL(None)
+        dlsym = libdl.dlsym
+        dlsym.restype = ctypes.c_void_p
+        dlsym.argtypes = [ctypes.c_void_p, ctypes.c_char_p]
+
+        stream_fn = dlsym(handle, b"_ZN2at3mps19getDefaultMPSStreamEv")
+        if not stream_fn:
+            return 0
+        stream_ptr = ctypes.CFUNCTYPE(ctypes.c_void_p)(stream_fn)()
+
+        cb_fn = dlsym(handle, b"_ZN2at3mps9MPSStream13commandBufferEv")
+        if not cb_fn:
+            return 0
+        cb_ptr = ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.c_void_p)(cb_fn)(stream_ptr)
+
+        objc = ctypes.CDLL("/usr/lib/libobjc.A.dylib")
+        sel_reg = objc.sel_registerName
+        sel_reg.restype = ctypes.c_void_p
+        sel_reg.argtypes = [ctypes.c_char_p]
+        msg_send = objc.objc_msgSend
+        msg_send.restype = ctypes.c_void_p
+        msg_send.argtypes = [ctypes.c_void_p, ctypes.c_void_p]
+
+        queue_ptr = msg_send(cb_ptr, sel_reg(b"commandQueue"))
+        return queue_ptr or 0
+    except (OSError, AttributeError):
+        return 0
+
+
 ########################## init ##########################
 def init(
     *,
@@ -255,6 +291,21 @@ def init(
             random_seed=seed,
         )
 
+    # On Metal, share PyTorch MPS's command queue with Quadrants so that GPU-side ordering is automatic and the
+    # per-interop-point sync overhead (qd.sync / torch.mps.synchronize) is eliminated.
+    if backend == _gs_backend.metal and device.type == "mps":
+        mps_queue = _get_mps_command_queue()
+        if not mps_queue:
+            raise_exception(
+                "Failed to extract PyTorch MPS's Metal command queue. "
+                "This is required on Apple Metal for correct GPU synchronisation between Genesis and PyTorch. "
+                "Please ensure you are using a supported PyTorch version (>= 2.0)."
+            )
+        qd_init_kwargs.update(
+            external_metal_command_queue=mps_queue,
+            external_metal_command_queue_is_torch_queue=True,
+        )
+
     # init quadrants
     qd_debug = debug and (os.environ.get("QD_DEBUG") != "0")
     with redirect_stdout(_qd_outputs):

genesis/engine/solvers/kinematic_solver.py

@@ -801,8 +801,6 @@ def set_qpos(self, qpos, qs_idx=None, envs_idx=None, *, skip_forward=False):
                 assign_indexed_tensor(data, mask, qpos)
                 if mask and isinstance(mask[0], torch.Tensor):
                     envs_idx = mask[0].reshape((-1,))
-            if gs.backend == gs.metal:
-                torch.mps.synchronize()
         else:
             qpos, qs_idx, envs_idx = self._sanitize_io_variables(
                 qpos, qs_idx, self.n_qs, "qs_idx", envs_idx, skip_allocation=True
@@ -869,8 +867,6 @@ def set_dofs_velocity(self, velocity, dofs_idx=None, envs_idx=None, *, skip_forw
                     assign_indexed_tensor(vel, mask, velocity)
                 if mask and isinstance(mask[0], torch.Tensor):
                     envs_idx = mask[0].reshape((-1,))
-            if gs.backend == gs.metal:
-                torch.mps.synchronize()
             if not skip_forward and not isinstance(envs_idx, torch.Tensor):
                 envs_idx = self._scene._sanitize_envs_idx(envs_idx)
         else:

genesis/engine/solvers/rigid/collider/collider.py

@@ -581,8 +581,6 @@ def reset(self, envs_idx=None, *, cache_only: bool = True) -> None:
                 normal.zero_()
             else:
                 normal[:, envs_idx] = 0.0
-            if gs.backend == gs.metal:
-                torch.mps.synchronize()
             return
 
         envs_idx = self._solver._scene._sanitize_envs_idx(envs_idx)
@@ -636,8 +634,6 @@ def clear(self, envs_idx=None):
                 pos[:, envs_idx] = 0.0
                 normal[:, envs_idx] = 0.0
                 force[:, envs_idx] = 0.0
-            if gs.backend == gs.metal:
-                torch.mps.synchronize()
             return
 
         if not isinstance(envs_idx, torch.Tensor):

genesis/engine/solvers/rigid/constraint/solver.py

@@ -152,8 +152,6 @@ def reset(self, envs_idx=None):
             else:
                 is_warmstart[envs_idx] = False
                 qacc_ws[:, envs_idx] = 0.0
-            if gs.backend == gs.metal:
-                torch.mps.synchronize()
             return
 
         envs_idx = self._solver._scene._sanitize_envs_idx(envs_idx)
@@ -186,8 +184,6 @@ def clear(self, envs_idx=None):
                 assign_indexed_tensor(n_constraints_equality, env_mask, 0)
                 assign_indexed_tensor(n_constraints_frictionloss, env_mask, 0)
                 assign_indexed_tensor(qd_n_equalities, env_mask, n_eq)
-            if gs.backend == gs.metal:
-                torch.mps.synchronize()
             return
 
         if not isinstance(envs_idx, torch.Tensor):

genesis/engine/solvers/rigid/rigid_solver.py

@@ -1568,8 +1568,6 @@ def set_state(self, f, state, envs_idx=None, *, partial: bool = False) -> None:
                 mass_dst[envs_idx] = state.mass_shift[envs_idx]
                 if self.n_geoms:
                     fric_dst[envs_idx] = state.friction_ratio[envs_idx]
-            if gs.backend == gs.metal:
-                torch.mps.synchronize()
         else:
             envs_idx = self._scene._sanitize_envs_idx(envs_idx)
             kernel_set_zero(envs_idx, self._errno)
@@ -1696,8 +1694,6 @@ def set_base_links_pos(self, pos, links_idx=None, envs_idx=None, *, relative=Fal
                 target = data[:, link.q_start : link.q_start + 3]
             pos = broadcast_tensor(pos, gs.tc_float, target.shape)
             torch.where(envs_idx[:, None], pos, target, out=target)
-            if gs.backend == gs.metal:
-                torch.mps.synchronize()
         else:
             pos, links_idx, envs_idx = self._sanitize_io_variables(
                 pos, links_idx, self.n_links, "links_idx", envs_idx, (3,), skip_allocation=True
@@ -1795,8 +1791,6 @@ def set_base_links_quat(self, quat, links_idx=None, envs_idx=None, *, relative=F
                 target = data[:, link.q_start + 3 : link.q_start + 7]
             quat = broadcast_tensor(quat, gs.tc_float, target.shape)
             torch.where(envs_idx[:, None], quat, target, out=target)
-            if gs.backend == gs.metal:
-                torch.mps.synchronize()
         else:
             quat, links_idx, envs_idx = self._sanitize_io_variables(
                 quat, links_idx, self.n_links, "links_idx", envs_idx, (4,), skip_allocation=True
@@ -1915,8 +1909,6 @@ def set_links_inertia(self, ratio, links_idx=None, envs_idx=None):
             assign_indexed_tensor(mass_data, mask, mass_data[mask] * ratio_t)
             assign_indexed_tensor(inertial_i_data, mask, inertial_i_data[mask] * ratio_t[..., None, None])
             assign_indexed_tensor(invweight_data, mask, invweight_data[mask] / ratio_t[..., None])
-            if gs.backend == gs.metal:
-                torch.mps.synchronize()
             return
 
         ratio, links_idx, envs_idx = self._sanitize_io_variables(
@@ -1971,8 +1963,6 @@ def set_qpos(self, qpos, qs_idx=None, envs_idx=None, *, skip_forward=False):
                 errno[envs_idx] = 0
                 if mask and isinstance(mask[0], torch.Tensor):
                     envs_idx = mask[0].reshape((-1,))
-            if gs.backend == gs.metal:
-                torch.mps.synchronize()
         else:
             qpos, qs_idx, envs_idx = self._sanitize_io_variables(
                 qpos, qs_idx, self.n_qs, "qs_idx", envs_idx, skip_allocation=True
@@ -2116,8 +2106,6 @@ def _set_dofs_info(self, tensor_list, dofs_idx, name, envs_idx=None):
                 num_values = len(tensor_list)
                 for j, mask_j in enumerate(((*mask, ..., j) for j in range(num_values)) if num_values > 1 else (mask,)):
                     assign_indexed_tensor(data, mask_j, tensor_list[j])
-            if gs.backend == gs.metal:
-                torch.mps.synchronize()
             return
 
         tensor_list = list(tensor_list)
@@ -2219,8 +2207,6 @@ def set_dofs_position(self, position, dofs_idx=None, envs_idx=None):
         if gs.use_zerocopy:
             errno = qd_to_torch(self._errno, copy=False)
             errno[envs_idx] = 0
-            if gs.backend == gs.metal:
-                torch.mps.synchronize()
         else:
             kernel_set_zero(envs_idx, self._errno)
 
@@ -2248,8 +2234,6 @@ def control_dofs_force(self, force, dofs_idx=None, envs_idx=None):
             ctrl_mode[mask] = gs.CTRL_MODE.FORCE
             ctrl_force = qd_to_torch(self.dofs_state.ctrl_force, transpose=True, copy=False)
             assign_indexed_tensor(ctrl_force, mask, force)
-            if gs.backend == gs.metal:
-                torch.mps.synchronize()
             return
 
         force, dofs_idx, envs_idx = self._sanitize_io_variables(
@@ -2269,8 +2253,6 @@ def control_dofs_velocity(self, velocity, dofs_idx=None, envs_idx=None):
             ctrl_pos[mask] = 0.0
             ctrl_vel = qd_to_torch(self.dofs_state.ctrl_vel, transpose=True, copy=False)
             assign_indexed_tensor(ctrl_vel, mask, velocity)
-            if gs.backend == gs.metal:
-                torch.mps.synchronize()
             return
 
         velocity, dofs_idx, envs_idx = self._sanitize_io_variables(
@@ -2290,8 +2272,6 @@ def control_dofs_position(self, position, dofs_idx=None, envs_idx=None):
             assign_indexed_tensor(ctrl_pos, mask, position)
             ctrl_vel = qd_to_torch(self.dofs_state.ctrl_vel, transpose=True, copy=False)
             ctrl_vel[mask] = 0.0
-            if gs.backend == gs.metal:
-                torch.mps.synchronize()
             return
 
         position, dofs_idx, envs_idx = self._sanitize_io_variables(
@@ -2311,8 +2291,6 @@ def control_dofs_position_velocity(self, position, velocity, dofs_idx=None, envs
             assign_indexed_tensor(ctrl_pos, mask, position)
             ctrl_vel = qd_to_torch(self.dofs_state.ctrl_vel, transpose=True, copy=False)
             assign_indexed_tensor(ctrl_vel, mask, velocity)
-            if gs.backend == gs.metal:
-                torch.mps.synchronize()
             return
 
         position, dofs_idx, _ = self._sanitize_io_variables(
@@ -2648,8 +2626,6 @@ def clear_external_force(self):
             for tensor in (self.links_state.cfrc_applied_ang, self.links_state.cfrc_applied_vel):
                 out = qd_to_torch(tensor, copy=False)
                 out.zero_()
-            if gs.backend == gs.metal:
-                torch.mps.synchronize()
             return
 
         kernel_clear_external_force(self.links_state, self._rigid_global_info, self._static_rigid_sim_config)

genesis/utils/misc.py

@@ -534,14 +534,6 @@ def qd_to_torch(
     else:
         try:
             tensor = value._T_tc if transpose else value._tc
-            # FIXME: DLPack may return old values on Apple Metal if sync is not systematically called manually.
-            # See comment in `qd_to_python` for details on the MPS command queue race condition.
-            if gs.backend == gs.metal:
-                qd.sync()
-            if copy:
-                tensor = tensor.clone()
-                if gs.backend == gs.metal:
-                    torch.mps.synchronize()
         except AttributeError:
             try:
                 tc = value.to_torch(copy=False)