Use full wave for all kernels except monolith

Author: doru1004

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

@@ -1845,7 +1845,7 @@ def func_hessian_and_cholesky_factor_direct(
 
     if qd.static(static_rigid_sim_config.backend == gs.cpu or static_rigid_sim_config.sparse_solve):
         # CPU
-        qd.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL, block_dim=32)
+        qd.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL, block_dim=qd.static(64 if static_rigid_sim_config.backend == gs.amdgpu else 32))
         for i_b in range(_B):
             func_hessian_and_cholesky_factor_direct_batch(
                 i_b,
@@ -1861,7 +1861,7 @@ def func_hessian_and_cholesky_factor_direct(
         if qd.static(static_rigid_sim_config.enable_tiled_cholesky_hessian):
             func_cholesky_factor_direct_tiled(constraint_state, rigid_global_info, static_rigid_sim_config)
         else:
-            qd.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL, block_dim=32)
+            qd.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL, block_dim=qd.static(64 if static_rigid_sim_config.backend == gs.amdgpu else 32))
             for i_b in range(_B):
                 func_cholesky_factor_direct_batch(i_b, constraint_state, rigid_global_info)
 
@@ -3246,7 +3246,7 @@ def func_update_gradient_tiled(
         )
 
     if qd.static(static_rigid_sim_config.solver_type == gs.constraint_solver.CG):
-        qd.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL, block_dim=32)
+        qd.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL, block_dim=qd.static(64 if static_rigid_sim_config.backend == gs.amdgpu else 32))
         for i_b in range(_B):
             func_solve_mass_batch(
                 i_b,
@@ -3289,7 +3289,7 @@ def func_update_gradient(
         not static_rigid_sim_config.enable_tiled_cholesky_hessian or static_rigid_sim_config.backend == gs.cpu
     ):
         # CPU
-        qd.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL, block_dim=32)
+        qd.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL, block_dim=qd.static(64 if static_rigid_sim_config.backend == gs.amdgpu else 32))
         for i_b in range(_B):
             func_update_gradient_batch(
                 i_b,