diff --git a/genesis/engine/solvers/rigid/collider/broadphase.py b/genesis/engine/solvers/rigid/collider/broadphase.py index d651baba46..928e7dc374 100644 --- a/genesis/engine/solvers/rigid/collider/broadphase.py +++ b/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 diff --git a/genesis/engine/solvers/rigid/constraint/solver.py b/genesis/engine/solvers/rigid/constraint/solver.py index cb7fbebed0..c947df9f9c 100644 --- a/genesis/engine/solvers/rigid/constraint/solver.py +++ b/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): diff --git a/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py b/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py index 43c15bd797..48c657e70c 100644 --- a/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py +++ b/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((ENVS, TWC_LDS_MAX_CON), gs.qd_float) # 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[env_in_block, i_c] = 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[env_in_block, j_c_lds] + 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