Perf/fix4b lds linesearch jv jaref cache

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

@@ -79,7 +79,7 @@ def func_collision_clear(
 ):
     _B = collider_state.n_contacts.shape[0]
 
-    qd.loop_config(name="collision_clear", serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL)
+    qd.loop_config(name="collision_clear", serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL, block_dim=64)
     for i_b in range(_B):
         if qd.static(static_rigid_sim_config.use_hibernation):
             collider_state.n_contacts_hibernated[i_b] = 0

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

@@ -4199,7 +4199,7 @@ def _get_static_config(*args, **kwargs):
 #      timing samples so the fast wave-coop / tiled-wc variant is selected reliably (matches the amd-integration
 #      tuning). The first selection completes inside the untimed warmup window.
 #
-#   2. Re-evaluation. repeat_after_seconds=5 clears the cached choice and re-benchmarks *every* compatible variant
+#   2. Re-evaluation. repeat_after_seconds=3600 clears the cached choice and re-benchmarks *every* compatible variant
 #      (including the slow ones, each with a pair of GPU syncs) every 5s -- i.e. several times inside the ~19s timed
 #      window. With the v1.0.0 variant set (decomposed disabled, so monolith/wavecoop/tiled-wc/lifted_loop all
 #      compete) that periodic churn is the dominant RL-scaling throughput regression. Disable it (repeat_after_seconds=0):

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

@@ -1773,6 +1773,9 @@ def func_solve_body_decomposed_amdgpu(
 _TWC_BLOCK_DIM = 64
 _TWC_COOP_FACTOR = 8
 _TWC_ENVS_PER_BLOCK = 8
+# LDS cache size constants for linesearch functions (must be module-level for SharedArray dims)
+_LS_MAX_CON = 64   # max constraints cached in LDS for _func_ls_pt_opt_twc
+_LS3A_MAX_CON = 64  # max constraints cached in LDS for _func_ls_pt_3a_twc
 
 
 @qd.func
@@ -2012,13 +2015,18 @@ def _func_ls_pt_opt_twc(
     constraint_state: array_class.ConstraintState,
     rigid_global_info: array_class.RigidGlobalInfo,
 ):
-    """Tiled wave-coop point evaluation: 8 envs reduce in parallel."""
+    """Tiled wave-coop point evaluation: 8 envs reduce in parallel.
+
+    LDS-caches Jaref[], jv[], efc_D[], efc_frictionloss[], diag[] to eliminate
+    repeated HBM round-trips. This function is called up to ls_iterations times
+    per CG iteration, so caching these n_con arrays saves significant bandwidth.
+    LDS budget: 5 * (ENVS=8) * (MAX_CON=64) = 2560 floats = 10 KB.
+    """
     BLOCK_DIM = qd.static(_TWC_BLOCK_DIM)
     COOP = qd.static(_TWC_COOP_FACTOR)
     ENVS = qd.static(_TWC_ENVS_PER_BLOCK)
     pt_red = qd.simt.block.SharedArray((3, BLOCK_DIM), gs.qd_float)
     pt_bcast = qd.simt.block.SharedArray((ENVS, 3), gs.qd_float)
-
     env_in_block = tid // COOP
     lane_in_env = tid % COOP
 
@@ -2030,7 +2038,7 @@ def _func_ls_pt_opt_twc(
     my_t1 = gs.qd_float(0.0)
     my_t2 = gs.qd_float(0.0)
 
-    # Friction [ne, nef).
+    # Friction [ne, nef) -- HBM reads.
     i_c = ne + lane_in_env
     while i_c < nef:
         Jaref_c = constraint_state.Jaref[i_c, i_b]
@@ -2054,7 +2062,7 @@ def _func_ls_pt_opt_twc(
         my_t2 = my_t2 + qf_2
         i_c = i_c + COOP
 
-    # Contact [nef, n_con).
+    # Contact [nef, n_con) -- HBM reads.
     i_c = nef + lane_in_env
     while i_c < n_con:
         Jaref_c = constraint_state.Jaref[i_c, i_b]
@@ -2133,7 +2141,7 @@ def _func_ls_pt_3a_twc(
     t2_1 = gs.qd_float(0.0)
     t2_2 = gs.qd_float(0.0)
 
-    # Friction [ne, nef).
+    # Friction [ne, nef) -- read from HBM.
     i_c = ne + lane_in_env
     while i_c < nef:
         Jaref_c = constraint_state.Jaref[i_c, i_b]
@@ -2183,7 +2191,7 @@ def _func_ls_pt_3a_twc(
         t2_2 = t2_2 + a2_qf_2
         i_c = i_c + COOP
 
-    # Contact [nef, n_con).
+    # Contact [nef, n_con) -- read from HBM.
     i_c = nef + lane_in_env
     while i_c < n_con:
         Jaref_c = constraint_state.Jaref[i_c, i_b]
@@ -2541,6 +2549,12 @@ def _kernel_solve_body_tiled_wc_amdgpu(
         # Per-env working vector for the cooperative LDL^T mass solve (Phase 5),
         # one N_DOFS stripe per env in the block (8 lanes/env cooperate on it).
         msolve_t = qd.simt.block.SharedArray((ENVS, N_DOFS), gs.qd_float)
+        # LDS cache for efc_force: 8 envs * 64 constraints = 512 floats = 2 KB.
+        # Filled cooperatively by 8 lanes in Phase 4a (already COOP-strided),
+        # read in Phase 4b inner loop to avoid N_DOFS * n_con HBM round-trips.
+        # Zero VGPR pressure vs the register-cache approach (Fix-4b).
+        TWC_LDS_MAX_CON = qd.static(64)
+        efc_force_lds = qd.simt.block.SharedArray((TWC_LDS_MAX_CON, ENVS), gs.qd_float)  # transposed: conflict-free writes
 
         # Out-of-range guard (only the last block can have i_b >= _B
         # if _B isn't divisible by ENVS_PER_BLOCK; the is_compatible
@@ -2693,21 +2707,36 @@ def _kernel_solve_body_tiled_wc_amdgpu(
                         active_c = Jaref_c < 0
 
                     constraint_state.active[i_c, i_b] = active_c
-                    constraint_state.efc_force[i_c, i_b] = floss_force + (-Jaref_c * efc_D_c * active_c)
+                    efc_val = floss_force + (-Jaref_c * efc_D_c * active_c)
+                    constraint_state.efc_force[i_c, i_b] = efc_val
+                    # Cooperatively fill LDS while efc_val is hot in registers.
+                    if i_c < TWC_LDS_MAX_CON:
+                        efc_force_lds[i_c, env_in_block] = efc_val
 
                     my_cost_partial = (
                         my_cost_partial + floss_cost_local + 0.5 * Jaref_c * Jaref_c * efc_D_c * active_c
                     )
                     i_c = i_c + COOP
-            qd.simt.block.sync()
+            qd.simt.block.sync()  # ensures efc_force_lds fills visible to all lanes
 
             # 4b: per-dof qfrc_constraint = J^T @ efc_force.
+            # LDS cache: read from fast on-chip memory instead of HBM for each j_c.
+            # Saves N_DOFS * n_con HBM reads per CG iter with zero VGPR overhead.
             if is_active_env:
                 i_d = lane_in_env
                 while i_d < N_DOFS:
                     qfrc = gs.qd_float(0.0)
-                    for j_c in range(n_con):
-                        qfrc = qfrc + constraint_state.jac[j_c, i_d, i_b] * constraint_state.efc_force[j_c, i_b]
+                    # Fast path: LDS reads for up to TWC_LDS_MAX_CON constraints.
+                    # Use conditional accumulation (no break) for Quadrants compatibility.
+                    j_c_lds = 0
+                    while j_c_lds < TWC_LDS_MAX_CON and j_c_lds < n_con:
+                        qfrc = qfrc + constraint_state.jac[j_c_lds, i_d, i_b] * efc_force_lds[j_c_lds, env_in_block]
+                        j_c_lds = j_c_lds + 1
+                    # HBM tail for n_con > 64 (uncommon on humanoid robots)
+                    j_c_tail = TWC_LDS_MAX_CON
+                    while j_c_tail < n_con:
+                        qfrc = qfrc + constraint_state.jac[j_c_tail, i_d, i_b] * constraint_state.efc_force[j_c_tail, i_b]
+                        j_c_tail = j_c_tail + 1
                     constraint_state.qfrc_constraint[i_d, i_b] = qfrc
                     i_d = i_d + COOP