[FEATURE] Avoid manually sync between torch and quadrants on Metal. (#2760)

Author: hughperkins

genesis/__init__.py

@@ -255,6 +255,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 = qd.interop.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

@@ -564,14 +564,8 @@ def qd_to_torch(
                 tensor = value._T_tc if transpose else value._tc
                 is_copy = False
 
-    if not is_copy:
-        # FIXME: DLPack may return old values on Apple Metal if sync is not systematically called manually
-        if gs.backend == gs.metal:
-            qd.sync()
-        if copy:
-            tensor = tensor.clone()
-            if gs.backend == gs.metal:
-                torch.mps.synchronize()
+    if copy and not is_copy:
+        tensor = tensor.clone()
 
     return _apply_masks(tensor, value, row_mask, col_mask, keepdim, copy, to_torch=True)