removed kernel for saving adjoint cache

Author: SonSang

genesis/engine/solvers/rigid/rigid_solver_decomp.py

@@ -921,16 +921,8 @@ def substep(self, f):
 
         self._links_state_cache.clear()
 
-        if f == 0 and self._requires_grad:
-            kernel_save_adjoint_cache(
-                f=f,
-                dofs_state=self.dofs_state,
-                rigid_global_info=self._rigid_global_info,
-                rigid_adjoint_cache=self._rigid_adjoint_cache,
-                static_rigid_sim_config=self._static_rigid_sim_config,
-            )
-
         kernel_step_1(
+            f=f,
             links_state=self.links_state,
             links_info=self.links_info,
             joints_state=self.joints_state,
@@ -942,6 +934,7 @@ def substep(self, f):
             entities_state=self.entities_state,
             entities_info=self.entities_info,
             rigid_global_info=self._rigid_global_info,
+            rigid_adjoint_cache=self._rigid_adjoint_cache,
             static_rigid_sim_config=self._static_rigid_sim_config,
             contact_island_state=self.constraint_solver.contact_island.contact_island_state,
             is_backward=False,
@@ -957,6 +950,7 @@ def substep(self, f):
         else:
             self._func_constraint_force()
             kernel_step_2(
+                f=f,
                 dofs_state=self.dofs_state,
                 dofs_info=self.dofs_info,
                 links_info=self.links_info,
@@ -969,20 +963,12 @@ def substep(self, f):
                 geoms_state=self.geoms_state,
                 collider_state=self.collider._collider_state,
                 rigid_global_info=self._rigid_global_info,
+                rigid_adjoint_cache=self._rigid_adjoint_cache,
                 static_rigid_sim_config=self._static_rigid_sim_config,
                 contact_island_state=self.constraint_solver.contact_island.contact_island_state,
                 is_backward=False,
             )
 
-            if self._requires_grad:
-                kernel_save_adjoint_cache(
-                    f=f + 1,
-                    dofs_state=self.dofs_state,
-                    rigid_global_info=self._rigid_global_info,
-                    rigid_adjoint_cache=self._rigid_adjoint_cache,
-                    static_rigid_sim_config=self._static_rigid_sim_config,
-                )
-
     def check_errno(self):
         # Note that errno must be evaluated BEFORE match because otherwise it will be evaluated for each case...
         # See official documentation: https://docs.python.org/3.10/reference/compound_stmts.html#overview
@@ -1245,8 +1231,6 @@ def substep_pre_coupling(self, f):
             self.substep(f)
 
     def substep_pre_coupling_grad(self, f):
-        curr_dofs_acc = self._rigid_adjoint_cache.dofs_acc.to_numpy()[f]
-
         # Run forward substep again to restore this step's information, this is needed because we do not store info
         # of every substep.
         kernel_prepare_backward_substep(
@@ -1311,6 +1295,7 @@ def substep_pre_coupling_grad(self, f):
             gs.raise_exception(f"Nan grad in qpos or dofs_vel found at step {self._sim.cur_step_global}")
 
         kernel_step_2.grad(
+            f=f,
             dofs_state=self.dofs_state,
             dofs_info=self.dofs_info,
             links_info=self.links_info,
@@ -1323,21 +1308,34 @@ def substep_pre_coupling_grad(self, f):
             geoms_state=self.geoms_state,
             collider_state=self.collider._collider_state,
             rigid_global_info=self._rigid_global_info,
+            rigid_adjoint_cache=self._rigid_adjoint_cache,
             static_rigid_sim_config=self._static_rigid_sim_config,
             contact_island_state=self.constraint_solver.contact_island.contact_island_state,
             is_backward=True,
         )
 
+        # We cannot use [kernel_forward_dynamics.grad] because we read [dofs_state.acc] and overwrite it in the kernel,
+        # which is prohibited (https://docs.taichi-lang.org/docs/differentiable_programming#global-data-access-rules).
+        # In [kernel_forward_dynamics], we read [acc] in [func_update_acc] and overwrite it in [kernel_compute_qacc].
+        # As [kenrel_compute_qacc] is called at the end of [kernel_forward_dynamics], we first backpropagate through
+        # [kernel_compute_qacc] and then restore the original [acc] from the adjoint cache. This copy operation
+        # cannot be merged with [kernel_compute_qacc.grad] because .grad function itself is a standalone kernel.
+        # We could possibly merge this small kernel later if (1) .grad function is regarded as a function instead of a
+        # kernel, (2) we add another variable to store the new [acc] from [kernel_compute_qacc] and thus can avoid
+        # the data access violation. However, both of these require major changes.
         kernel_compute_qacc.grad(
             dofs_state=self.dofs_state,
             entities_info=self.entities_info,
             rigid_global_info=self._rigid_global_info,
             static_rigid_sim_config=self._static_rigid_sim_config,
             is_backward=True,
         )
-
-        # Load the current dofs_acc from adjoint cache, as it was overwritten by [kernel_compute_qacc]
-        self.dofs_state.acc.from_numpy(curr_dofs_acc)
+        kernel_copy_acc(
+            f=f,
+            dofs_state=self.dofs_state,
+            rigid_adjoint_cache=self._rigid_adjoint_cache,
+            static_rigid_sim_config=self._static_rigid_sim_config,
+        )
 
         kernel_forward_dynamics_without_qacc.grad(
             links_state=self.links_state,
@@ -1386,6 +1384,7 @@ def substep_post_coupling(self, f):
                 is_backward=False,
             )
             kernel_step_2(
+                f=f,
                 dofs_state=self.dofs_state,
                 dofs_info=self.dofs_info,
                 links_info=self.links_info,
@@ -1398,18 +1397,11 @@ def substep_post_coupling(self, f):
                 geoms_state=self.geoms_state,
                 collider_state=self.collider._collider_state,
                 rigid_global_info=self._rigid_global_info,
+                rigid_adjoint_cache=self._rigid_adjoint_cache,
                 static_rigid_sim_config=self._static_rigid_sim_config,
                 contact_island_state=self.constraint_solver.contact_island.contact_island_state,
                 is_backward=False,
             )
-            if self._requires_grad:
-                kernel_save_adjoint_cache(
-                    f=f + 1,
-                    dofs_state=self.dofs_state,
-                    rigid_global_info=self._rigid_global_info,
-                    rigid_adjoint_cache=self._rigid_adjoint_cache,
-                    static_rigid_sim_config=self._static_rigid_sim_config,
-                )
         elif isinstance(self.sim.coupler, IPCCoupler):
             # For IPCCoupler, perform full rigid body computation in post-coupling phase
             # This allows IPC to handle rigid bodies during the coupling phase
@@ -4574,6 +4566,7 @@ def func_update_cartesian_space(
 
 @ti.kernel(fastcache=gs.use_fastcache)
 def kernel_step_1(
+    f: ti.int32,
     links_state: array_class.LinksState,
     links_info: array_class.LinksInfo,
     joints_state: array_class.JointsState,
@@ -4585,10 +4578,21 @@ def kernel_step_1(
     entities_state: array_class.EntitiesState,
     entities_info: array_class.EntitiesInfo,
     rigid_global_info: array_class.RigidGlobalInfo,
+    rigid_adjoint_cache: array_class.RigidAdjointCache,
     static_rigid_sim_config: ti.template(),
     contact_island_state: array_class.ContactIslandState,
     is_backward: ti.template(),
 ):
+    if ti.static(static_rigid_sim_config.requires_grad):
+        if f == 0:
+            func_save_adjoint_cache(
+                f=f,
+                dofs_state=dofs_state,
+                rigid_global_info=rigid_global_info,
+                rigid_adjoint_cache=rigid_adjoint_cache,
+                static_rigid_sim_config=static_rigid_sim_config,
+            )
+
     if ti.static(static_rigid_sim_config.enable_mujoco_compatibility):
         _B = links_state.pos.shape[1]
         ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL)
@@ -4700,6 +4704,7 @@ def func_implicit_damping(
 
 @ti.kernel(fastcache=gs.use_fastcache)
 def kernel_step_2(
+    f: ti.int32,
     dofs_state: array_class.DofsState,
     dofs_info: array_class.DofsInfo,
     links_info: array_class.LinksInfo,
@@ -4712,6 +4717,7 @@ def kernel_step_2(
     geoms_state: array_class.GeomsState,
     collider_state: array_class.ColliderState,
     rigid_global_info: array_class.RigidGlobalInfo,
+    rigid_adjoint_cache: array_class.RigidAdjointCache,
     static_rigid_sim_config: ti.template(),
     contact_island_state: array_class.ContactIslandState,
     is_backward: ti.template(),
@@ -4799,6 +4805,15 @@ def kernel_step_2(
                     is_backward=is_backward,
                 )
 
+        if ti.static(static_rigid_sim_config.requires_grad):
+            func_save_adjoint_cache(
+                f=f + 1,
+                dofs_state=dofs_state,
+                rigid_global_info=rigid_global_info,
+                rigid_adjoint_cache=rigid_adjoint_cache,
+                static_rigid_sim_config=static_rigid_sim_config,
+            )
+
 
 @ti.kernel(fastcache=gs.use_fastcache)
 def kernel_forward_kinematics_links_geoms(
@@ -6753,19 +6768,6 @@ def func_integrate(
                                 )
 
 
-@ti.kernel(fastcache=gs.use_fastcache)
-def kernel_copy_next_to_curr(
-    dofs_state: array_class.DofsState,
-    rigid_global_info: array_class.RigidGlobalInfo,
-    static_rigid_sim_config: ti.template(),
-):
-    func_copy_next_to_curr(
-        dofs_state=dofs_state,
-        rigid_global_info=rigid_global_info,
-        static_rigid_sim_config=static_rigid_sim_config,
-    )
-
-
 @ti.func
 def func_copy_next_to_curr(
     dofs_state: array_class.DofsState,
@@ -6806,8 +6808,8 @@ def func_copy_next_to_curr_grad(
         rigid_global_info.qpos[i_q, i_b] = rigid_adjoint_cache.qpos[f, i_q, i_b]
 
 
-@ti.kernel(fastcache=gs.use_fastcache)
-def kernel_save_adjoint_cache(
+@ti.func
+def func_save_adjoint_cache(
     f: ti.int32,
     dofs_state: array_class.DofsState,
     rigid_global_info: array_class.RigidGlobalInfo,
@@ -7048,28 +7050,18 @@ def func_copy_cartesian_space(
 
 
 @ti.kernel(fastcache=gs.use_fastcache)
-def kernel_copy_cartesian_space(
-    src_dofs_state: array_class.DofsState,
-    src_links_state: array_class.LinksState,
-    src_joints_state: array_class.JointsState,
-    src_geoms_state: array_class.GeomsState,
-    dst_dofs_state: array_class.DofsState,
-    dst_links_state: array_class.LinksState,
-    dst_joints_state: array_class.JointsState,
-    dst_geoms_state: array_class.GeomsState,
+def kernel_copy_acc(
+    f: ti.int32,
+    dofs_state: array_class.DofsState,
+    rigid_adjoint_cache: array_class.RigidAdjointCache,
     static_rigid_sim_config: ti.template(),
 ):
-    func_copy_cartesian_space(
-        src_dofs_state=src_dofs_state,
-        src_links_state=src_links_state,
-        src_joints_state=src_joints_state,
-        src_geoms_state=src_geoms_state,
-        dst_dofs_state=dst_dofs_state,
-        dst_links_state=dst_links_state,
-        dst_joints_state=dst_joints_state,
-        dst_geoms_state=dst_geoms_state,
-        static_rigid_sim_config=static_rigid_sim_config,
-    )
+    n_dofs = dofs_state.vel.shape[0]
+    _B = dofs_state.vel.shape[1]
+
+    ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL)
+    for i_d, i_b in ti.ndrange(n_dofs, _B):
+        dofs_state.acc[i_d, i_b] = rigid_adjoint_cache.dofs_acc[f, i_d, i_b]
 
 
 @ti.func

tests/test_grad.py

@@ -193,6 +193,7 @@ def constraint_solver_resolve():
     # Step once to compute constraint solver's inputs: [mass], [jac], [aref], [efc_D], [force]. We do not call the
     # entire scene.step() because it will overwrite the necessary information that we need to compute the gradients.
     kernel_step_1(
+        f=0,
         links_state=rigid_solver.links_state,
         links_info=rigid_solver.links_info,
         joints_state=rigid_solver.joints_state,
@@ -204,6 +205,7 @@ def constraint_solver_resolve():
         entities_state=rigid_solver.entities_state,
         entities_info=rigid_solver.entities_info,
         rigid_global_info=rigid_solver._rigid_global_info,
+        rigid_adjoint_cache=rigid_solver._rigid_adjoint_cache,
         static_rigid_sim_config=rigid_solver._static_rigid_sim_config,
         contact_island_state=constraint_solver.contact_island.contact_island_state,
         is_backward=False,