From 93f659a7a836f3b17913ab7508f891425e505601 Mon Sep 17 00:00:00 2001 From: Zhihui Du Date: Wed, 1 Jul 2026 19:43:30 -0700 Subject: [PATCH 1/6] 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 From 6469f5ae10fd54b53b1d5de425a0f7e03e98e2ac Mon Sep 17 00:00:00 2001 From: Zhihui Du Date: Sat, 18 Jul 2026 15:03:15 -0400 Subject: [PATCH 2/6] fix: increase perf_dispatch warmup/active so LDS variant wins benchmark correctly --- genesis/engine/solvers/rigid/constraint/solver.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/genesis/engine/solvers/rigid/constraint/solver.py b/genesis/engine/solvers/rigid/constraint/solver.py index cb7fbebed..f7c9fc682 100644 --- a/genesis/engine/solvers/rigid/constraint/solver.py +++ b/genesis/engine/solvers/rigid/constraint/solver.py @@ -4207,9 +4207,9 @@ def _get_static_config(*args, **kwargs): # window with zero dispatch overhead (the chosen impl is served from the cached fast path). @qd.perf_dispatch( get_geometry_hash=lambda *args, **kwargs: (*args, frozendict(kwargs)), - first_warmup=3, - warmup=3, - active=5, + first_warmup=10, + warmup=10, + active=15, repeat_after_seconds=0, ) def func_solve_body( From 77a6429e58b8bf4f4be7628944784081ccdfced0 Mon Sep 17 00:00:00 2001 From: Zhihui Du Date: Sat, 18 Jul 2026 15:19:50 -0400 Subject: [PATCH 3/6] revert: restore perf_dispatch params to baseline (first_warmup=3,warmup=3,active=5) --- genesis/engine/solvers/rigid/constraint/solver.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/genesis/engine/solvers/rigid/constraint/solver.py b/genesis/engine/solvers/rigid/constraint/solver.py index f7c9fc682..cb7fbebed 100644 --- a/genesis/engine/solvers/rigid/constraint/solver.py +++ b/genesis/engine/solvers/rigid/constraint/solver.py @@ -4207,9 +4207,9 @@ def _get_static_config(*args, **kwargs): # window with zero dispatch overhead (the chosen impl is served from the cached fast path). @qd.perf_dispatch( get_geometry_hash=lambda *args, **kwargs: (*args, frozendict(kwargs)), - first_warmup=10, - warmup=10, - active=15, + first_warmup=3, + warmup=3, + active=5, repeat_after_seconds=0, ) def func_solve_body( From 986971e875ca553d7022131b9b970c393bfda95f Mon Sep 17 00:00:00 2001 From: Zhihui Du Date: Sat, 18 Jul 2026 17:08:48 -0400 Subject: [PATCH 4/6] fix: transpose efc_force_lds to [MAX_CON,ENVS] to eliminate 8-way LDS bank conflicts on Phase 4a write --- genesis/engine/solvers/rigid/constraint/solver_amdgpu.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py b/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py index f604142cd..baf36ed33 100644 --- a/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py +++ b/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py @@ -2546,7 +2546,7 @@ def _kernel_solve_body_tiled_wc_amdgpu( # 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) + 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 @@ -2703,7 +2703,7 @@ def _kernel_solve_body_tiled_wc_amdgpu( 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 + 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 @@ -2722,7 +2722,7 @@ def _kernel_solve_body_tiled_wc_amdgpu( # 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] + 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 From 2b8bcd8337a902f7486f0a6b28a42249e42fea47 Mon Sep 17 00:00:00 2001 From: Zhihui Du Date: Sat, 18 Jul 2026 17:48:57 -0400 Subject: [PATCH 5/6] perf: increase efc_force_lds MAX_CON from 64 to 256 for full G1 constraint coverage --- genesis/engine/solvers/rigid/constraint/solver_amdgpu.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py b/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py index baf36ed33..1b7a01c00 100644 --- a/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py +++ b/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py @@ -2545,7 +2545,7 @@ def _kernel_solve_body_tiled_wc_amdgpu( # 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) + TWC_LDS_MAX_CON = qd.static(256) # increased: covers G1 ~200 constraints fully 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 From 8f46258b655b40768663e0e97fba75d31fb6a08e Mon Sep 17 00:00:00 2001 From: Zhihui Du Date: Sat, 18 Jul 2026 18:03:27 -0400 Subject: [PATCH 6/6] revert: MAX_CON back to 64 (256 reduces occupancy, net negative for current workloads) --- genesis/engine/solvers/rigid/constraint/solver_amdgpu.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py b/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py index 1b7a01c00..baf36ed33 100644 --- a/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py +++ b/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py @@ -2545,7 +2545,7 @@ def _kernel_solve_body_tiled_wc_amdgpu( # 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(256) # increased: covers G1 ~200 constraints fully + 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