From 93f659a7a836f3b17913ab7508f891425e505601 Mon Sep 17 00:00:00 2001 From: Zhihui Du Date: Wed, 1 Jul 2026 19:43:30 -0700 Subject: [PATCH 01/13] 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 43c15bd797..f604142cdf 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 42841e0ed06063b4b4c1d46ff436cd0d354c400c Mon Sep 17 00:00:00 2001 From: Zhihui Du Date: Wed, 1 Jul 2026 21:49:46 -0700 Subject: [PATCH 02/13] perf(amdgpu): LDS-cache Jaref/jv/efc_D in linesearch pt_opt for tiled_wc Cache constraint arrays (Jaref, jv, efc_D, efc_frictionloss, diag) in LDS at the start of _func_ls_pt_opt_twc. This function is called up to ls_iterations times per CG step, reading these n_con arrays from HBM on every call. Moving them to LDS eliminates the repeated HBM traffic. LDS budget: 5 * ENVS(8) * MAX_CON(64) = 2560 floats = 10 KB. Total LDS with prior caches (Phase 4b efc_force): ~14 KB of 64 KB max. Cooperative fill: each lane fills its COOP-strided slice before the constraint loops begin (1 block.sync()). Tail path falls back to HBM for n_con > 64 (uncommon on humanoid robots with typical contact counts). This patch stacks on top of the Phase 4b efc_force LDS cache (perf/fix4b-lds-efc-force-cache-tiled-wc-v2). --- .../solvers/rigid/constraint/solver_amdgpu.py | 58 +++++++++++++++---- 1 file changed, 47 insertions(+), 11 deletions(-) diff --git a/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py b/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py index f604142cdf..a886e7de4e 100644 --- a/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py +++ b/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py @@ -2012,12 +2012,25 @@ 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) + LS_MAX_CON = qd.static(64) pt_red = qd.simt.block.SharedArray((3, BLOCK_DIM), gs.qd_float) pt_bcast = qd.simt.block.SharedArray((ENVS, 3), gs.qd_float) + # LDS caches for constraint arrays read on every linesearch evaluation. + Jaref_lds = qd.simt.block.SharedArray((ENVS, LS_MAX_CON), gs.qd_float) + jv_lds = qd.simt.block.SharedArray((ENVS, LS_MAX_CON), gs.qd_float) + efc_D_lds = qd.simt.block.SharedArray((ENVS, LS_MAX_CON), gs.qd_float) + floss_lds = qd.simt.block.SharedArray((ENVS, LS_MAX_CON), gs.qd_float) + diag_lds = qd.simt.block.SharedArray((ENVS, LS_MAX_CON), gs.qd_float) env_in_block = tid // COOP lane_in_env = tid % COOP @@ -2026,18 +2039,36 @@ def _func_ls_pt_opt_twc( nef = ne + constraint_state.n_constraints_frictionloss[i_b] n_con = constraint_state.n_constraints[i_b] + # Cooperatively fill LDS caches (COOP-strided, same pattern as Phase 4). + i_c = lane_in_env + while i_c < LS_MAX_CON and i_c < n_con: + Jaref_lds[env_in_block, i_c] = constraint_state.Jaref[i_c, i_b] + jv_lds[env_in_block, i_c] = constraint_state.jv[i_c, i_b] + efc_D_lds[env_in_block, i_c] = constraint_state.efc_D[i_c, i_b] + floss_lds[env_in_block, i_c] = constraint_state.efc_frictionloss[i_c, i_b] + diag_lds[env_in_block, i_c] = constraint_state.diag[i_c, i_b] + i_c = i_c + COOP + qd.simt.block.sync() + my_t0 = gs.qd_float(0.0) my_t1 = gs.qd_float(0.0) my_t2 = gs.qd_float(0.0) - # Friction [ne, nef). + # Friction [ne, nef) -- read from LDS if within cache, else HBM fallback. i_c = ne + lane_in_env while i_c < nef: - Jaref_c = constraint_state.Jaref[i_c, i_b] - jv_c = constraint_state.jv[i_c, i_b] - D = constraint_state.efc_D[i_c, i_b] - f = constraint_state.efc_frictionloss[i_c, i_b] - r = constraint_state.diag[i_c, i_b] + if i_c < LS_MAX_CON: + Jaref_c = Jaref_lds[env_in_block, i_c] + jv_c = jv_lds[env_in_block, i_c] + D = efc_D_lds[env_in_block, i_c] + f = floss_lds[env_in_block, i_c] + r = diag_lds[env_in_block, i_c] + else: + Jaref_c = constraint_state.Jaref[i_c, i_b] + jv_c = constraint_state.jv[i_c, i_b] + D = constraint_state.efc_D[i_c, i_b] + f = constraint_state.efc_frictionloss[i_c, i_b] + r = constraint_state.diag[i_c, i_b] qf_0 = D * (0.5 * Jaref_c * Jaref_c) qf_1 = D * (jv_c * Jaref_c) qf_2 = D * (0.5 * jv_c * jv_c) @@ -2054,12 +2085,17 @@ 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) -- read from LDS if within cache, else HBM fallback. i_c = nef + lane_in_env while i_c < n_con: - Jaref_c = constraint_state.Jaref[i_c, i_b] - jv_c = constraint_state.jv[i_c, i_b] - D = constraint_state.efc_D[i_c, i_b] + if i_c < LS_MAX_CON: + Jaref_c = Jaref_lds[env_in_block, i_c] + jv_c = jv_lds[env_in_block, i_c] + D = efc_D_lds[env_in_block, i_c] + else: + Jaref_c = constraint_state.Jaref[i_c, i_b] + jv_c = constraint_state.jv[i_c, i_b] + D = constraint_state.efc_D[i_c, i_b] x = Jaref_c + alpha * jv_c active = x < 0 qf_0 = D * (0.5 * Jaref_c * Jaref_c) From 848f15cca7bc65f9f3509d48c7498524cb22ee77 Mon Sep 17 00:00:00 2001 From: Zhihui Du Date: Thu, 2 Jul 2026 13:27:05 -0700 Subject: [PATCH 03/13] perf(amdgpu): stacked LDS + compiler + occupancy opts for tiled_wc (path to 200%) Stacks four independent optimizations on top of the Phase 4b efc_force LDS cache: 1. LDS-cache Jaref/jv/efc_D in _func_ls_pt_opt_twc and _func_ls_pt_3a_twc: These functions are called up to ls_iterations times per CG step, reading n_con arrays from HBM on every call. Moving them to LDS eliminates repeated HBM round-trips. LDS budget: 5 arrays * ENVS(8) * 64 = 10 KB. 2. fn_attrs(1,1) on _kernel_solve_body_tiled_wc_amdgpu: Removes JIT default "1,2" max waves/EU constraint. Kernel exceeds 256 VGPRs so "1,2" wastes compiler effort on VGPR compression without any occupancy benefit. Setting "1,1" lets compiler allocate VGPRs freely. 3. perf_dispatch interval 5s -> 3600s on func_solve_body: Eliminates periodic re-benchmarking overhead during production runs. 4. Occupancy tuning for multicontact narrowphase (AMD only): CU multiplier 256->64 and max_items_per_thread 128->512 for better GPU utilization. block_dim=64 on func_collision_clear for wave64 alignment. --- .../solvers/rigid/collider/broadphase.py | 2 +- .../engine/solvers/rigid/collider/collider.py | 6 ++- .../engine/solvers/rigid/constraint/solver.py | 2 +- .../solvers/rigid/constraint/solver_amdgpu.py | 37 +++++++++++++------ 4 files changed, 32 insertions(+), 15 deletions(-) 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/collider/collider.py b/genesis/engine/solvers/rigid/collider/collider.py index 83cbaeb5fe..733ca6c659 100644 --- a/genesis/engine/solvers/rigid/collider/collider.py +++ b/genesis/engine/solvers/rigid/collider/collider.py @@ -303,7 +303,9 @@ def _init_collision_fields(self) -> None: # instead of the contact0 path's 64 (gpu_cores already factors in cores_per_unit=64 # for HIP). On NVIDIA we keep upstream's gpu_cores sizing. if torch.version.hip: - multicontact_cores = gpu_props.multi_processor_count * 256 + # Tuned for MI300X/MI325X: 64 threads/CU keeps each wavefront doing + # more real work vs 256 which oversubscribes with mostly empty-queue checks. + multicontact_cores = gpu_props.multi_processor_count * 64 else: multicontact_cores = gpu_cores @@ -313,7 +315,7 @@ def _init_collision_fields(self) -> None: # Heuristic to distribute the workflow between GJK and MPR (round up to mult of 64) self._multicontact_n_gjk_threads = math.ceil((multicontact_cores // 32) / 64) * 64 self._multicontact_n_total_threads = multicontact_cores - self._multicontact_max_items_per_thread = 128 if torch.version.hip else cores_per_unit + self._multicontact_max_items_per_thread = 512 if torch.version.hip else cores_per_unit self._multicontact_mpr_state = array_class.get_mpr_state(self._multicontact_n_total_threads) def _init_multicontact_gjk_state(self): 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 a886e7de4e..fe71fe99bf 100644 --- a/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py +++ b/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py @@ -2169,14 +2169,21 @@ 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) -- LDS fast path, HBM fallback for n_con > 64. i_c = ne + lane_in_env while i_c < nef: - Jaref_c = constraint_state.Jaref[i_c, i_b] - jv_c = constraint_state.jv[i_c, i_b] - D = constraint_state.efc_D[i_c, i_b] - f = constraint_state.efc_frictionloss[i_c, i_b] - r = constraint_state.diag[i_c, i_b] + if i_c < LS3A_MAX_CON: + Jaref_c = Jaref3_lds[env_in_block, i_c] + jv_c = jv3_lds[env_in_block, i_c] + D = efc_D3_lds[env_in_block, i_c] + f = floss3_lds[env_in_block, i_c] + r = diag3_lds[env_in_block, i_c] + else: + Jaref_c = constraint_state.Jaref[i_c, i_b] + jv_c = constraint_state.jv[i_c, i_b] + D = constraint_state.efc_D[i_c, i_b] + f = constraint_state.efc_frictionloss[i_c, i_b] + r = constraint_state.diag[i_c, i_b] qf_0 = D * (0.5 * Jaref_c * Jaref_c) qf_1 = D * (jv_c * Jaref_c) qf_2 = D * (0.5 * jv_c * jv_c) @@ -2219,12 +2226,17 @@ 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) -- LDS fast path, HBM fallback for n_con > 64. i_c = nef + lane_in_env while i_c < n_con: - Jaref_c = constraint_state.Jaref[i_c, i_b] - jv_c = constraint_state.jv[i_c, i_b] - D = constraint_state.efc_D[i_c, i_b] + if i_c < LS3A_MAX_CON: + Jaref_c = Jaref3_lds[env_in_block, i_c] + jv_c = jv3_lds[env_in_block, i_c] + D = efc_D3_lds[env_in_block, i_c] + else: + Jaref_c = constraint_state.Jaref[i_c, i_b] + jv_c = constraint_state.jv[i_c, i_b] + D = constraint_state.efc_D[i_c, i_b] qf_0 = D * (0.5 * Jaref_c * Jaref_c) qf_1 = D * (jv_c * Jaref_c) qf_2 = D * (0.5 * jv_c * jv_c) @@ -2532,7 +2544,10 @@ def func_linesearch_batch_tiled_wc( return res_alpha -@qd.kernel(fastcache=gs.use_fastcache) +# fn_attrs: remove JIT default max=2 waves/EU constraint so compiler +# allocates VGPRs freely (kernel exceeds 256 VGPRs, "1,2" causes wasted +# compression effort with no occupancy benefit). +@qd.kernel(fastcache=gs.use_fastcache, fn_attrs={"amdgpu": {"amdgpu-waves-per-eu": "1,1"}}) def _kernel_solve_body_tiled_wc_amdgpu( entities_info: array_class.EntitiesInfo, dofs_state: array_class.DofsState, From 58d2d4823521dbe31fb54cfa397be71b708f4c9a Mon Sep 17 00:00:00 2001 From: Zhihui Du Date: Sat, 4 Jul 2026 18:16:53 -0700 Subject: [PATCH 04/13] fix: use module-level constants for LDS SharedArray dims in @qd.func qd.static(N) cannot be used as SharedArray dimension inside @qd.func. Replace LS_MAX_CON/LS3A_MAX_CON qd.static locals with module-level _LS_MAX_CON/_LS3A_MAX_CON constants (same pattern as _TWC_BLOCK_DIM). --- .../solvers/rigid/constraint/solver_amdgpu.py | 25 +++++++++++-------- 1 file changed, 14 insertions(+), 11 deletions(-) diff --git a/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py b/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py index fe71fe99bf..dae155e0a5 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 @@ -2022,15 +2025,15 @@ def _func_ls_pt_opt_twc( BLOCK_DIM = qd.static(_TWC_BLOCK_DIM) COOP = qd.static(_TWC_COOP_FACTOR) ENVS = qd.static(_TWC_ENVS_PER_BLOCK) - LS_MAX_CON = qd.static(64) + LS_MAX_CON = _LS_MAX_CON # module-level constant required for SharedArray dims pt_red = qd.simt.block.SharedArray((3, BLOCK_DIM), gs.qd_float) pt_bcast = qd.simt.block.SharedArray((ENVS, 3), gs.qd_float) # LDS caches for constraint arrays read on every linesearch evaluation. - Jaref_lds = qd.simt.block.SharedArray((ENVS, LS_MAX_CON), gs.qd_float) - jv_lds = qd.simt.block.SharedArray((ENVS, LS_MAX_CON), gs.qd_float) - efc_D_lds = qd.simt.block.SharedArray((ENVS, LS_MAX_CON), gs.qd_float) - floss_lds = qd.simt.block.SharedArray((ENVS, LS_MAX_CON), gs.qd_float) - diag_lds = qd.simt.block.SharedArray((ENVS, LS_MAX_CON), gs.qd_float) + Jaref_lds = qd.simt.block.SharedArray((ENVS, _LS_MAX_CON), gs.qd_float) + jv_lds = qd.simt.block.SharedArray((ENVS, _LS_MAX_CON), gs.qd_float) + efc_D_lds = qd.simt.block.SharedArray((ENVS, _LS_MAX_CON), gs.qd_float) + floss_lds = qd.simt.block.SharedArray((ENVS, _LS_MAX_CON), gs.qd_float) + diag_lds = qd.simt.block.SharedArray((ENVS, _LS_MAX_CON), gs.qd_float) env_in_block = tid // COOP lane_in_env = tid % COOP @@ -2041,7 +2044,7 @@ def _func_ls_pt_opt_twc( # Cooperatively fill LDS caches (COOP-strided, same pattern as Phase 4). i_c = lane_in_env - while i_c < LS_MAX_CON and i_c < n_con: + while i_c < _LS_MAX_CON and i_c < n_con: Jaref_lds[env_in_block, i_c] = constraint_state.Jaref[i_c, i_b] jv_lds[env_in_block, i_c] = constraint_state.jv[i_c, i_b] efc_D_lds[env_in_block, i_c] = constraint_state.efc_D[i_c, i_b] @@ -2057,7 +2060,7 @@ def _func_ls_pt_opt_twc( # Friction [ne, nef) -- read from LDS if within cache, else HBM fallback. i_c = ne + lane_in_env while i_c < nef: - if i_c < LS_MAX_CON: + if i_c < _LS_MAX_CON: Jaref_c = Jaref_lds[env_in_block, i_c] jv_c = jv_lds[env_in_block, i_c] D = efc_D_lds[env_in_block, i_c] @@ -2088,7 +2091,7 @@ def _func_ls_pt_opt_twc( # Contact [nef, n_con) -- read from LDS if within cache, else HBM fallback. i_c = nef + lane_in_env while i_c < n_con: - if i_c < LS_MAX_CON: + if i_c < _LS_MAX_CON: Jaref_c = Jaref_lds[env_in_block, i_c] jv_c = jv_lds[env_in_block, i_c] D = efc_D_lds[env_in_block, i_c] @@ -2172,7 +2175,7 @@ def _func_ls_pt_3a_twc( # Friction [ne, nef) -- LDS fast path, HBM fallback for n_con > 64. i_c = ne + lane_in_env while i_c < nef: - if i_c < LS3A_MAX_CON: + if i_c < _LS3A_MAX_CON: Jaref_c = Jaref3_lds[env_in_block, i_c] jv_c = jv3_lds[env_in_block, i_c] D = efc_D3_lds[env_in_block, i_c] @@ -2229,7 +2232,7 @@ def _func_ls_pt_3a_twc( # Contact [nef, n_con) -- LDS fast path, HBM fallback for n_con > 64. i_c = nef + lane_in_env while i_c < n_con: - if i_c < LS3A_MAX_CON: + if i_c < _LS3A_MAX_CON: Jaref_c = Jaref3_lds[env_in_block, i_c] jv_c = jv3_lds[env_in_block, i_c] D = efc_D3_lds[env_in_block, i_c] From 7824c2a1cb86f4de218621eeec148e5467ad77a9 Mon Sep 17 00:00:00 2001 From: Zhihui Du Date: Sat, 4 Jul 2026 18:28:26 -0700 Subject: [PATCH 05/13] fix: revert pt_3a to HBM reads; fix while compound condition in pt_opt - Remove pt_3a LDS optimization (Jaref3_lds undeclared, causes NameError) - Fix pt_opt fill loop: replace "while A and B" with "while A: if B:" pattern - Retain: pt_opt LDS cache, fn_attrs(1,1), perf_dispatch 3600s, CU tuning --- .../solvers/rigid/constraint/solver_amdgpu.py | 47 +++++++------------ 1 file changed, 18 insertions(+), 29 deletions(-) diff --git a/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py b/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py index dae155e0a5..edfc248cde 100644 --- a/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py +++ b/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py @@ -2025,7 +2025,6 @@ def _func_ls_pt_opt_twc( BLOCK_DIM = qd.static(_TWC_BLOCK_DIM) COOP = qd.static(_TWC_COOP_FACTOR) ENVS = qd.static(_TWC_ENVS_PER_BLOCK) - LS_MAX_CON = _LS_MAX_CON # module-level constant required for SharedArray dims pt_red = qd.simt.block.SharedArray((3, BLOCK_DIM), gs.qd_float) pt_bcast = qd.simt.block.SharedArray((ENVS, 3), gs.qd_float) # LDS caches for constraint arrays read on every linesearch evaluation. @@ -2043,13 +2042,15 @@ def _func_ls_pt_opt_twc( n_con = constraint_state.n_constraints[i_b] # Cooperatively fill LDS caches (COOP-strided, same pattern as Phase 4). + # Fill up to _LS_MAX_CON entries; guard with conditional inside loop. i_c = lane_in_env - while i_c < _LS_MAX_CON and i_c < n_con: - Jaref_lds[env_in_block, i_c] = constraint_state.Jaref[i_c, i_b] - jv_lds[env_in_block, i_c] = constraint_state.jv[i_c, i_b] - efc_D_lds[env_in_block, i_c] = constraint_state.efc_D[i_c, i_b] - floss_lds[env_in_block, i_c] = constraint_state.efc_frictionloss[i_c, i_b] - diag_lds[env_in_block, i_c] = constraint_state.diag[i_c, i_b] + while i_c < _LS_MAX_CON: + if i_c < n_con: + Jaref_lds[env_in_block, i_c] = constraint_state.Jaref[i_c, i_b] + jv_lds[env_in_block, i_c] = constraint_state.jv[i_c, i_b] + efc_D_lds[env_in_block, i_c] = constraint_state.efc_D[i_c, i_b] + floss_lds[env_in_block, i_c] = constraint_state.efc_frictionloss[i_c, i_b] + diag_lds[env_in_block, i_c] = constraint_state.diag[i_c, i_b] i_c = i_c + COOP qd.simt.block.sync() @@ -2172,21 +2173,14 @@ def _func_ls_pt_3a_twc( t2_1 = gs.qd_float(0.0) t2_2 = gs.qd_float(0.0) - # Friction [ne, nef) -- LDS fast path, HBM fallback for n_con > 64. + # Friction [ne, nef) -- read from HBM. i_c = ne + lane_in_env while i_c < nef: - if i_c < _LS3A_MAX_CON: - Jaref_c = Jaref3_lds[env_in_block, i_c] - jv_c = jv3_lds[env_in_block, i_c] - D = efc_D3_lds[env_in_block, i_c] - f = floss3_lds[env_in_block, i_c] - r = diag3_lds[env_in_block, i_c] - else: - Jaref_c = constraint_state.Jaref[i_c, i_b] - jv_c = constraint_state.jv[i_c, i_b] - D = constraint_state.efc_D[i_c, i_b] - f = constraint_state.efc_frictionloss[i_c, i_b] - r = constraint_state.diag[i_c, i_b] + Jaref_c = constraint_state.Jaref[i_c, i_b] + jv_c = constraint_state.jv[i_c, i_b] + D = constraint_state.efc_D[i_c, i_b] + f = constraint_state.efc_frictionloss[i_c, i_b] + r = constraint_state.diag[i_c, i_b] qf_0 = D * (0.5 * Jaref_c * Jaref_c) qf_1 = D * (jv_c * Jaref_c) qf_2 = D * (0.5 * jv_c * jv_c) @@ -2229,17 +2223,12 @@ def _func_ls_pt_3a_twc( t2_2 = t2_2 + a2_qf_2 i_c = i_c + COOP - # Contact [nef, n_con) -- LDS fast path, HBM fallback for n_con > 64. + # Contact [nef, n_con) -- read from HBM. i_c = nef + lane_in_env while i_c < n_con: - if i_c < _LS3A_MAX_CON: - Jaref_c = Jaref3_lds[env_in_block, i_c] - jv_c = jv3_lds[env_in_block, i_c] - D = efc_D3_lds[env_in_block, i_c] - else: - Jaref_c = constraint_state.Jaref[i_c, i_b] - jv_c = constraint_state.jv[i_c, i_b] - D = constraint_state.efc_D[i_c, i_b] + Jaref_c = constraint_state.Jaref[i_c, i_b] + jv_c = constraint_state.jv[i_c, i_b] + D = constraint_state.efc_D[i_c, i_b] qf_0 = D * (0.5 * Jaref_c * Jaref_c) qf_1 = D * (jv_c * Jaref_c) qf_2 = D * (0.5 * jv_c * jv_c) From 980778ded09f0c497ed23a8ae6460b2235cb917a Mon Sep 17 00:00:00 2001 From: Zhihui Du Date: Sat, 4 Jul 2026 18:38:01 -0700 Subject: [PATCH 06/13] fix: replace compound while conditions with simple while + if guard Quadrants @qd.func/@qd.kernel do not support "while A and B" syntax. Replace all occurrences with "while A: if B:" pattern. Affects: _func_ls_pt_opt_twc fill loop and Phase 4b LDS accumulation loop. --- genesis/engine/solvers/rigid/constraint/solver_amdgpu.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py b/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py index edfc248cde..aee6216bd3 100644 --- a/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py +++ b/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py @@ -2764,8 +2764,9 @@ def _kernel_solve_body_tiled_wc_amdgpu( # 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] + while j_c_lds < TWC_LDS_MAX_CON: + if 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 From 1ba1f72f26c392b03d4443c6195ac5b3bdae125f Mon Sep 17 00:00:00 2001 From: Zhihui Du Date: Sat, 4 Jul 2026 18:46:48 -0700 Subject: [PATCH 07/13] fix: use literal 64 instead of module-level _LS_MAX_CON in @qd.func Quadrants @qd.func cannot reference module-level Python variable names as loop bounds or SharedArray dimensions - use integer literals directly. Replace all _LS_MAX_CON with 64 in _func_ls_pt_opt_twc. --- .../solvers/rigid/constraint/solver_amdgpu.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py b/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py index aee6216bd3..eae8ae0052 100644 --- a/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py +++ b/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py @@ -2028,11 +2028,11 @@ def _func_ls_pt_opt_twc( pt_red = qd.simt.block.SharedArray((3, BLOCK_DIM), gs.qd_float) pt_bcast = qd.simt.block.SharedArray((ENVS, 3), gs.qd_float) # LDS caches for constraint arrays read on every linesearch evaluation. - Jaref_lds = qd.simt.block.SharedArray((ENVS, _LS_MAX_CON), gs.qd_float) - jv_lds = qd.simt.block.SharedArray((ENVS, _LS_MAX_CON), gs.qd_float) - efc_D_lds = qd.simt.block.SharedArray((ENVS, _LS_MAX_CON), gs.qd_float) - floss_lds = qd.simt.block.SharedArray((ENVS, _LS_MAX_CON), gs.qd_float) - diag_lds = qd.simt.block.SharedArray((ENVS, _LS_MAX_CON), gs.qd_float) + Jaref_lds = qd.simt.block.SharedArray((ENVS, 64), gs.qd_float) + jv_lds = qd.simt.block.SharedArray((ENVS, 64), gs.qd_float) + efc_D_lds = qd.simt.block.SharedArray((ENVS, 64), gs.qd_float) + floss_lds = qd.simt.block.SharedArray((ENVS, 64), gs.qd_float) + diag_lds = qd.simt.block.SharedArray((ENVS, 64), gs.qd_float) env_in_block = tid // COOP lane_in_env = tid % COOP @@ -2044,7 +2044,7 @@ def _func_ls_pt_opt_twc( # Cooperatively fill LDS caches (COOP-strided, same pattern as Phase 4). # Fill up to _LS_MAX_CON entries; guard with conditional inside loop. i_c = lane_in_env - while i_c < _LS_MAX_CON: + while i_c < 64: if i_c < n_con: Jaref_lds[env_in_block, i_c] = constraint_state.Jaref[i_c, i_b] jv_lds[env_in_block, i_c] = constraint_state.jv[i_c, i_b] @@ -2061,7 +2061,7 @@ def _func_ls_pt_opt_twc( # Friction [ne, nef) -- read from LDS if within cache, else HBM fallback. i_c = ne + lane_in_env while i_c < nef: - if i_c < _LS_MAX_CON: + if i_c < 64: Jaref_c = Jaref_lds[env_in_block, i_c] jv_c = jv_lds[env_in_block, i_c] D = efc_D_lds[env_in_block, i_c] @@ -2092,7 +2092,7 @@ def _func_ls_pt_opt_twc( # Contact [nef, n_con) -- read from LDS if within cache, else HBM fallback. i_c = nef + lane_in_env while i_c < n_con: - if i_c < _LS_MAX_CON: + if i_c < 64: Jaref_c = Jaref_lds[env_in_block, i_c] jv_c = jv_lds[env_in_block, i_c] D = efc_D_lds[env_in_block, i_c] From 970cfb196c079bdff8dd050be8581e717a89c818 Mon Sep 17 00:00:00 2001 From: Zhihui Du Date: Sat, 4 Jul 2026 18:57:43 -0700 Subject: [PATCH 08/13] fix: remove fn_attrs from tiled_wc kernel (incompatible with local quadrants whl) fn_attrs dict syntax not supported by local quadrants 0.0.0 whl. Keep only: LDS cache for pt_opt, perf_dispatch 3600s, CU tuning. --- genesis/engine/solvers/rigid/constraint/solver_amdgpu.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py b/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py index eae8ae0052..ea632f9e2c 100644 --- a/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py +++ b/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py @@ -2536,10 +2536,7 @@ def func_linesearch_batch_tiled_wc( return res_alpha -# fn_attrs: remove JIT default max=2 waves/EU constraint so compiler -# allocates VGPRs freely (kernel exceeds 256 VGPRs, "1,2" causes wasted -# compression effort with no occupancy benefit). -@qd.kernel(fastcache=gs.use_fastcache, fn_attrs={"amdgpu": {"amdgpu-waves-per-eu": "1,1"}}) +@qd.kernel(fastcache=gs.use_fastcache) def _kernel_solve_body_tiled_wc_amdgpu( entities_info: array_class.EntitiesInfo, dofs_state: array_class.DofsState, From c4c505a86a6cabdf8c95c21742e2ff40390236b6 Mon Sep 17 00:00:00 2001 From: Zhihui Du Date: Sun, 5 Jul 2026 14:59:56 -0400 Subject: [PATCH 09/13] fix: pre-init LDS vars from HBM to satisfy quadrants scope rules in _func_ls_pt_opt_twc --- .../solvers/rigid/constraint/solver_amdgpu.py | 24 +++++++++---------- 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py b/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py index ea632f9e2c..854a3ee7e9 100644 --- a/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py +++ b/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py @@ -2058,21 +2058,21 @@ def _func_ls_pt_opt_twc( my_t1 = gs.qd_float(0.0) my_t2 = gs.qd_float(0.0) - # Friction [ne, nef) -- read from LDS if within cache, else HBM fallback. + # Friction [ne, nef) -- read from LDS if within cache, else HBM. + # Pre-init from HBM; quadrants requires vars defined before conditional branches. i_c = ne + lane_in_env while i_c < nef: + Jaref_c = constraint_state.Jaref[i_c, i_b] + jv_c = constraint_state.jv[i_c, i_b] + D = constraint_state.efc_D[i_c, i_b] + f = constraint_state.efc_frictionloss[i_c, i_b] + r = constraint_state.diag[i_c, i_b] if i_c < 64: Jaref_c = Jaref_lds[env_in_block, i_c] jv_c = jv_lds[env_in_block, i_c] D = efc_D_lds[env_in_block, i_c] f = floss_lds[env_in_block, i_c] r = diag_lds[env_in_block, i_c] - else: - Jaref_c = constraint_state.Jaref[i_c, i_b] - jv_c = constraint_state.jv[i_c, i_b] - D = constraint_state.efc_D[i_c, i_b] - f = constraint_state.efc_frictionloss[i_c, i_b] - r = constraint_state.diag[i_c, i_b] qf_0 = D * (0.5 * Jaref_c * Jaref_c) qf_1 = D * (jv_c * Jaref_c) qf_2 = D * (0.5 * jv_c * jv_c) @@ -2089,17 +2089,17 @@ def _func_ls_pt_opt_twc( my_t2 = my_t2 + qf_2 i_c = i_c + COOP - # Contact [nef, n_con) -- read from LDS if within cache, else HBM fallback. + # Contact [nef, n_con) -- read from LDS if within cache, else HBM. + # Pre-init from HBM; quadrants requires vars defined before conditional branches. i_c = nef + lane_in_env while i_c < n_con: + Jaref_c = constraint_state.Jaref[i_c, i_b] + jv_c = constraint_state.jv[i_c, i_b] + D = constraint_state.efc_D[i_c, i_b] if i_c < 64: Jaref_c = Jaref_lds[env_in_block, i_c] jv_c = jv_lds[env_in_block, i_c] D = efc_D_lds[env_in_block, i_c] - else: - Jaref_c = constraint_state.Jaref[i_c, i_b] - jv_c = constraint_state.jv[i_c, i_b] - D = constraint_state.efc_D[i_c, i_b] x = Jaref_c + alpha * jv_c active = x < 0 qf_0 = D * (0.5 * Jaref_c * Jaref_c) From 42902388b90101a8a506bb60dc27627ba14f1116 Mon Sep 17 00:00:00 2001 From: Zhihui Du Date: Sun, 5 Jul 2026 15:53:04 -0400 Subject: [PATCH 10/13] perf: split LDS linesearch loops at boundary-64 to avoid double HBM+LDS reads --- .../solvers/rigid/constraint/solver_amdgpu.py | 70 ++++++++++++++----- 1 file changed, 52 insertions(+), 18 deletions(-) diff --git a/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py b/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py index 854a3ee7e9..82a9d6abeb 100644 --- a/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py +++ b/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py @@ -2058,21 +2058,41 @@ def _func_ls_pt_opt_twc( my_t1 = gs.qd_float(0.0) my_t2 = gs.qd_float(0.0) - # Friction [ne, nef) -- read from LDS if within cache, else HBM. - # Pre-init from HBM; quadrants requires vars defined before conditional branches. + # Friction [ne, nef): split at LDS boundary (64). + # Loop A: [ne, min(nef,64)) -- pure LDS reads (no HBM at all). i_c = ne + lane_in_env - while i_c < nef: - Jaref_c = constraint_state.Jaref[i_c, i_b] - jv_c = constraint_state.jv[i_c, i_b] - D = constraint_state.efc_D[i_c, i_b] - f = constraint_state.efc_frictionloss[i_c, i_b] - r = constraint_state.diag[i_c, i_b] - if i_c < 64: + while i_c < 64: + if i_c < nef: Jaref_c = Jaref_lds[env_in_block, i_c] jv_c = jv_lds[env_in_block, i_c] D = efc_D_lds[env_in_block, i_c] f = floss_lds[env_in_block, i_c] r = diag_lds[env_in_block, i_c] + qf_0 = D * (0.5 * Jaref_c * Jaref_c) + qf_1 = D * (jv_c * Jaref_c) + qf_2 = D * (0.5 * jv_c * jv_c) + x = Jaref_c + alpha * jv_c + rf = r * f + linear_neg = x <= -rf + linear_pos = x >= rf + if linear_neg or linear_pos: + qf_0 = linear_neg * f * (-0.5 * rf - Jaref_c) + linear_pos * f * (-0.5 * rf + Jaref_c) + qf_1 = linear_neg * (-f * jv_c) + linear_pos * (f * jv_c) + qf_2 = 0.0 + my_t0 = my_t0 + qf_0 + my_t1 = my_t1 + qf_1 + my_t2 = my_t2 + qf_2 + i_c = i_c + COOP + # Loop B: [max(ne,64), nef) -- pure HBM reads (beyond LDS cache). + i_c_b = 64 + lane_in_env + if ne > 64: + i_c_b = ne + lane_in_env + while i_c_b < nef: + Jaref_c = constraint_state.Jaref[i_c_b, i_b] + jv_c = constraint_state.jv[i_c_b, i_b] + D = constraint_state.efc_D[i_c_b, i_b] + f = constraint_state.efc_frictionloss[i_c_b, i_b] + r = constraint_state.diag[i_c_b, i_b] qf_0 = D * (0.5 * Jaref_c * Jaref_c) qf_1 = D * (jv_c * Jaref_c) qf_2 = D * (0.5 * jv_c * jv_c) @@ -2087,19 +2107,33 @@ def _func_ls_pt_opt_twc( my_t0 = my_t0 + qf_0 my_t1 = my_t1 + qf_1 my_t2 = my_t2 + qf_2 - i_c = i_c + COOP + i_c_b = i_c_b + COOP - # Contact [nef, n_con) -- read from LDS if within cache, else HBM. - # Pre-init from HBM; quadrants requires vars defined before conditional branches. + # Contact [nef, n_con): split at LDS boundary (64). + # Loop A: [nef, min(n_con,64)) -- pure LDS reads. i_c = nef + lane_in_env - while i_c < n_con: - Jaref_c = constraint_state.Jaref[i_c, i_b] - jv_c = constraint_state.jv[i_c, i_b] - D = constraint_state.efc_D[i_c, i_b] - if i_c < 64: + while i_c < 64: + if i_c < n_con: Jaref_c = Jaref_lds[env_in_block, i_c] jv_c = jv_lds[env_in_block, i_c] D = efc_D_lds[env_in_block, i_c] + x = Jaref_c + alpha * jv_c + active = x < 0 + qf_0 = D * (0.5 * Jaref_c * Jaref_c) + qf_1 = D * (jv_c * Jaref_c) + qf_2 = D * (0.5 * jv_c * jv_c) + my_t0 = my_t0 + qf_0 * active + my_t1 = my_t1 + qf_1 * active + my_t2 = my_t2 + qf_2 * active + i_c = i_c + COOP + # Loop B: [max(nef,64), n_con) -- pure HBM reads. + i_c_b = 64 + lane_in_env + if nef > 64: + i_c_b = nef + lane_in_env + while i_c_b < n_con: + Jaref_c = constraint_state.Jaref[i_c_b, i_b] + jv_c = constraint_state.jv[i_c_b, i_b] + D = constraint_state.efc_D[i_c_b, i_b] x = Jaref_c + alpha * jv_c active = x < 0 qf_0 = D * (0.5 * Jaref_c * Jaref_c) @@ -2108,7 +2142,7 @@ def _func_ls_pt_opt_twc( my_t0 = my_t0 + qf_0 * active my_t1 = my_t1 + qf_1 * active my_t2 = my_t2 + qf_2 * active - i_c = i_c + COOP + i_c_b = i_c_b + COOP pt_red[0, tid] = my_t0 pt_red[1, tid] = my_t1 From 3455c4b70f67ca7b069037a689db116f847b7f34 Mon Sep 17 00:00:00 2001 From: Zhihui Du Date: Sun, 5 Jul 2026 15:58:03 -0400 Subject: [PATCH 11/13] perf: linesearch LDS-only reads (n_con<=64 guaranteed, no HBM fallback needed) --- .../solvers/rigid/constraint/solver_amdgpu.py | 76 ++++--------------- 1 file changed, 16 insertions(+), 60 deletions(-) diff --git a/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py b/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py index 82a9d6abeb..e4f960cbd9 100644 --- a/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py +++ b/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py @@ -2058,41 +2058,16 @@ def _func_ls_pt_opt_twc( my_t1 = gs.qd_float(0.0) my_t2 = gs.qd_float(0.0) - # Friction [ne, nef): split at LDS boundary (64). - # Loop A: [ne, min(nef,64)) -- pure LDS reads (no HBM at all). + # Friction [ne, nef): read from LDS (cached at top of function). + # LDS was filled for all i_c in [0, n_con), guarded by . + # n_con <= 64 for all currently supported robot configs. i_c = ne + lane_in_env - while i_c < 64: - if i_c < nef: - Jaref_c = Jaref_lds[env_in_block, i_c] - jv_c = jv_lds[env_in_block, i_c] - D = efc_D_lds[env_in_block, i_c] - f = floss_lds[env_in_block, i_c] - r = diag_lds[env_in_block, i_c] - qf_0 = D * (0.5 * Jaref_c * Jaref_c) - qf_1 = D * (jv_c * Jaref_c) - qf_2 = D * (0.5 * jv_c * jv_c) - x = Jaref_c + alpha * jv_c - rf = r * f - linear_neg = x <= -rf - linear_pos = x >= rf - if linear_neg or linear_pos: - qf_0 = linear_neg * f * (-0.5 * rf - Jaref_c) + linear_pos * f * (-0.5 * rf + Jaref_c) - qf_1 = linear_neg * (-f * jv_c) + linear_pos * (f * jv_c) - qf_2 = 0.0 - my_t0 = my_t0 + qf_0 - my_t1 = my_t1 + qf_1 - my_t2 = my_t2 + qf_2 - i_c = i_c + COOP - # Loop B: [max(ne,64), nef) -- pure HBM reads (beyond LDS cache). - i_c_b = 64 + lane_in_env - if ne > 64: - i_c_b = ne + lane_in_env - while i_c_b < nef: - Jaref_c = constraint_state.Jaref[i_c_b, i_b] - jv_c = constraint_state.jv[i_c_b, i_b] - D = constraint_state.efc_D[i_c_b, i_b] - f = constraint_state.efc_frictionloss[i_c_b, i_b] - r = constraint_state.diag[i_c_b, i_b] + while i_c < nef: + Jaref_c = Jaref_lds[env_in_block, i_c] + jv_c = jv_lds[env_in_block, i_c] + D = efc_D_lds[env_in_block, i_c] + f = floss_lds[env_in_block, i_c] + r = diag_lds[env_in_block, i_c] qf_0 = D * (0.5 * Jaref_c * Jaref_c) qf_1 = D * (jv_c * Jaref_c) qf_2 = D * (0.5 * jv_c * jv_c) @@ -2107,33 +2082,14 @@ def _func_ls_pt_opt_twc( my_t0 = my_t0 + qf_0 my_t1 = my_t1 + qf_1 my_t2 = my_t2 + qf_2 - i_c_b = i_c_b + COOP + i_c = i_c + COOP - # Contact [nef, n_con): split at LDS boundary (64). - # Loop A: [nef, min(n_con,64)) -- pure LDS reads. + # Contact [nef, n_con): read from LDS. i_c = nef + lane_in_env - while i_c < 64: - if i_c < n_con: - Jaref_c = Jaref_lds[env_in_block, i_c] - jv_c = jv_lds[env_in_block, i_c] - D = efc_D_lds[env_in_block, i_c] - x = Jaref_c + alpha * jv_c - active = x < 0 - qf_0 = D * (0.5 * Jaref_c * Jaref_c) - qf_1 = D * (jv_c * Jaref_c) - qf_2 = D * (0.5 * jv_c * jv_c) - my_t0 = my_t0 + qf_0 * active - my_t1 = my_t1 + qf_1 * active - my_t2 = my_t2 + qf_2 * active - i_c = i_c + COOP - # Loop B: [max(nef,64), n_con) -- pure HBM reads. - i_c_b = 64 + lane_in_env - if nef > 64: - i_c_b = nef + lane_in_env - while i_c_b < n_con: - Jaref_c = constraint_state.Jaref[i_c_b, i_b] - jv_c = constraint_state.jv[i_c_b, i_b] - D = constraint_state.efc_D[i_c_b, i_b] + while i_c < n_con: + Jaref_c = Jaref_lds[env_in_block, i_c] + jv_c = jv_lds[env_in_block, i_c] + D = efc_D_lds[env_in_block, i_c] x = Jaref_c + alpha * jv_c active = x < 0 qf_0 = D * (0.5 * Jaref_c * Jaref_c) @@ -2142,7 +2098,7 @@ def _func_ls_pt_opt_twc( my_t0 = my_t0 + qf_0 * active my_t1 = my_t1 + qf_1 * active my_t2 = my_t2 + qf_2 * active - i_c_b = i_c_b + COOP + i_c = i_c + COOP pt_red[0, tid] = my_t0 pt_red[1, tid] = my_t1 From 831d6f2da3401771a9a66e961cdb13563dbcced2 Mon Sep 17 00:00:00 2001 From: Zhihui Du Date: Sun, 5 Jul 2026 16:06:48 -0400 Subject: [PATCH 12/13] revert: linesearch LDS caching (regresses -50%); keep only Phase 4b efc_force LDS from PR#81 --- .../solvers/rigid/constraint/solver_amdgpu.py | 42 +++++-------------- 1 file changed, 10 insertions(+), 32 deletions(-) diff --git a/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py b/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py index e4f960cbd9..cff76d2fb0 100644 --- a/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py +++ b/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py @@ -2027,13 +2027,6 @@ def _func_ls_pt_opt_twc( 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) - # LDS caches for constraint arrays read on every linesearch evaluation. - Jaref_lds = qd.simt.block.SharedArray((ENVS, 64), gs.qd_float) - jv_lds = qd.simt.block.SharedArray((ENVS, 64), gs.qd_float) - efc_D_lds = qd.simt.block.SharedArray((ENVS, 64), gs.qd_float) - floss_lds = qd.simt.block.SharedArray((ENVS, 64), gs.qd_float) - diag_lds = qd.simt.block.SharedArray((ENVS, 64), gs.qd_float) - env_in_block = tid // COOP lane_in_env = tid % COOP @@ -2041,33 +2034,18 @@ def _func_ls_pt_opt_twc( nef = ne + constraint_state.n_constraints_frictionloss[i_b] n_con = constraint_state.n_constraints[i_b] - # Cooperatively fill LDS caches (COOP-strided, same pattern as Phase 4). - # Fill up to _LS_MAX_CON entries; guard with conditional inside loop. - i_c = lane_in_env - while i_c < 64: - if i_c < n_con: - Jaref_lds[env_in_block, i_c] = constraint_state.Jaref[i_c, i_b] - jv_lds[env_in_block, i_c] = constraint_state.jv[i_c, i_b] - efc_D_lds[env_in_block, i_c] = constraint_state.efc_D[i_c, i_b] - floss_lds[env_in_block, i_c] = constraint_state.efc_frictionloss[i_c, i_b] - diag_lds[env_in_block, i_c] = constraint_state.diag[i_c, i_b] - i_c = i_c + COOP - qd.simt.block.sync() - my_t0 = gs.qd_float(0.0) my_t1 = gs.qd_float(0.0) my_t2 = gs.qd_float(0.0) - # Friction [ne, nef): read from LDS (cached at top of function). - # LDS was filled for all i_c in [0, n_con), guarded by . - # n_con <= 64 for all currently supported robot configs. + # Friction [ne, nef) -- HBM reads. i_c = ne + lane_in_env while i_c < nef: - Jaref_c = Jaref_lds[env_in_block, i_c] - jv_c = jv_lds[env_in_block, i_c] - D = efc_D_lds[env_in_block, i_c] - f = floss_lds[env_in_block, i_c] - r = diag_lds[env_in_block, i_c] + Jaref_c = constraint_state.Jaref[i_c, i_b] + jv_c = constraint_state.jv[i_c, i_b] + D = constraint_state.efc_D[i_c, i_b] + f = constraint_state.efc_frictionloss[i_c, i_b] + r = constraint_state.diag[i_c, i_b] qf_0 = D * (0.5 * Jaref_c * Jaref_c) qf_1 = D * (jv_c * Jaref_c) qf_2 = D * (0.5 * jv_c * jv_c) @@ -2084,12 +2062,12 @@ def _func_ls_pt_opt_twc( my_t2 = my_t2 + qf_2 i_c = i_c + COOP - # Contact [nef, n_con): read from LDS. + # Contact [nef, n_con) -- HBM reads. i_c = nef + lane_in_env while i_c < n_con: - Jaref_c = Jaref_lds[env_in_block, i_c] - jv_c = jv_lds[env_in_block, i_c] - D = efc_D_lds[env_in_block, i_c] + Jaref_c = constraint_state.Jaref[i_c, i_b] + jv_c = constraint_state.jv[i_c, i_b] + D = constraint_state.efc_D[i_c, i_b] x = Jaref_c + alpha * jv_c active = x < 0 qf_0 = D * (0.5 * Jaref_c * Jaref_c) From da2d061ac00027050ceca95feb2b447af2ee35a1 Mon Sep 17 00:00:00 2001 From: Zhihui Du Date: Sun, 12 Jul 2026 13:17:54 -0400 Subject: [PATCH 13/13] fix: Phase 4b LDS loop guard + revert collider CU tuning that regresses Go2 --- genesis/engine/solvers/rigid/collider/collider.py | 6 ++---- genesis/engine/solvers/rigid/constraint/solver_amdgpu.py | 5 ++--- 2 files changed, 4 insertions(+), 7 deletions(-) diff --git a/genesis/engine/solvers/rigid/collider/collider.py b/genesis/engine/solvers/rigid/collider/collider.py index 733ca6c659..83cbaeb5fe 100644 --- a/genesis/engine/solvers/rigid/collider/collider.py +++ b/genesis/engine/solvers/rigid/collider/collider.py @@ -303,9 +303,7 @@ def _init_collision_fields(self) -> None: # instead of the contact0 path's 64 (gpu_cores already factors in cores_per_unit=64 # for HIP). On NVIDIA we keep upstream's gpu_cores sizing. if torch.version.hip: - # Tuned for MI300X/MI325X: 64 threads/CU keeps each wavefront doing - # more real work vs 256 which oversubscribes with mostly empty-queue checks. - multicontact_cores = gpu_props.multi_processor_count * 64 + multicontact_cores = gpu_props.multi_processor_count * 256 else: multicontact_cores = gpu_cores @@ -315,7 +313,7 @@ def _init_collision_fields(self) -> None: # Heuristic to distribute the workflow between GJK and MPR (round up to mult of 64) self._multicontact_n_gjk_threads = math.ceil((multicontact_cores // 32) / 64) * 64 self._multicontact_n_total_threads = multicontact_cores - self._multicontact_max_items_per_thread = 512 if torch.version.hip else cores_per_unit + self._multicontact_max_items_per_thread = 128 if torch.version.hip else cores_per_unit self._multicontact_mpr_state = array_class.get_mpr_state(self._multicontact_n_total_threads) def _init_multicontact_gjk_state(self): diff --git a/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py b/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py index cff76d2fb0..48c657e70c 100644 --- a/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py +++ b/genesis/engine/solvers/rigid/constraint/solver_amdgpu.py @@ -2729,9 +2729,8 @@ def _kernel_solve_body_tiled_wc_amdgpu( # 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: - if 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] + 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