From 93f659a7a836f3b17913ab7508f891425e505601 Mon Sep 17 00:00:00 2001 From: Zhihui Du Date: Wed, 1 Jul 2026 19:43:30 -0700 Subject: [PATCH] perf(amdgpu): LDS-cache efc_force in tiled_wc Phase 4b Replace per-lane HBM reads of efc_force in the J^T@efc_force inner loop with cooperative LDS fills during Phase 4a, then fast LDS reads in Phase 4b. The 8 lanes per env already write efc_force to HBM in COOP-strided order during Phase 4a. This patch stores efc_val into efc_force_lds at the same time (while the value is hot in registers), then Phase 4b reads from LDS instead of HBM for the inner j_c accumulation loop. LDS budget: (ENVS=8, MAX_CON=64) float32 = 2 KB, well within the 64 KB per-workgroup LDS limit on gfx942. No VGPR overhead vs the HBM path. Why LDS over register cache (Fix-4b): - Register cache consumed 8*8=64 VGPRs per lane, hurting occupancy on MI325X - LDS cache uses shared on-chip memory with zero VGPR cost - LDS latency (~100 cycles) vs HBM (~600 cycles) still gives significant speedup - Tail path handles n_con > 64 via HBM fallback (uncommon on humanoid robots) Correctness: efc_force_lds is filled before the existing block.sync() that Phase 4b already depends on, so no additional synchronization is needed. --- .../solvers/rigid/constraint/solver_amdgpu.py | 29 ++++++++++++++++--- 1 file changed, 25 insertions(+), 4 deletions(-) diff --git a/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py b/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py index 43c15bd79..f604142cd 100644 --- a/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py +++ b/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py @@ -2541,6 +2541,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 +2699,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