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

Author: duburcqa

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

@@ -838,8 +838,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
@@ -906,8 +904,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

@@ -585,8 +585,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)
@@ -640,8 +638,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

@@ -1592,8 +1592,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)
@@ -1720,8 +1718,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
@@ -1819,8 +1815,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
@@ -1939,8 +1933,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(
@@ -1995,8 +1987,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
@@ -2140,8 +2130,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)
@@ -2243,8 +2231,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)
 
@@ -2272,8 +2258,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(
@@ -2293,8 +2277,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(
@@ -2314,8 +2296,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(
@@ -2335,8 +2315,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(
@@ -2672,8 +2650,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

@@ -568,14 +568,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)