restore cache_kernel (#1318)

Author: thowell

mujoco_warp/_src/collision_convex.py

@@ -42,6 +42,7 @@
 from mujoco_warp._src.types import mat43
 from mujoco_warp._src.types import mat63
 from mujoco_warp._src.types import vec5
+from mujoco_warp._src.warp_util import cache_kernel
 from mujoco_warp._src.warp_util import event_scope
 
 # TODO(team): improve compile time to enable backward pass
@@ -153,6 +154,7 @@ def _hfield_filter(
     return False, xmin, xmax, ymin, ymax, zmin, zmax
 
 
+@cache_kernel
 def ccd_hfield_kernel_builder(
   geomtype1: int,
   geomtype2: int,
@@ -695,6 +697,7 @@ def ccd_hfield_kernel(
   return ccd_hfield_kernel
 
 
+@cache_kernel
 def ccd_kernel_builder(
   geomtype1: int,
   geomtype2: int,

mujoco_warp/_src/collision_driver.py

@@ -34,6 +34,7 @@
 from mujoco_warp._src.types import Model
 from mujoco_warp._src.types import mat23
 from mujoco_warp._src.types import mat63
+from mujoco_warp._src.warp_util import cache_kernel
 from mujoco_warp._src.warp_util import event_scope
 
 wp.set_module_options({"enable_backward": False})
@@ -441,6 +442,7 @@ def _sap_range(
   range_out[worldid, geomid] = limit - geomid
 
 
+@cache_kernel
 def _sap_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: int, ngeom_margin: int):
   @wp.kernel(module="unique", enable_backward=False)
   def kernel(
@@ -641,6 +643,7 @@ def sap_broadphase(m: Model, d: Data, ctx: CollisionContext):
   )
 
 
+@cache_kernel
 def _nxn_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: int, ngeom_margin: int):
   @wp.kernel(module="unique", enable_backward=False)
   def kernel(

mujoco_warp/_src/collision_primitive.py

@@ -42,6 +42,7 @@
 from mujoco_warp._src.types import Model
 from mujoco_warp._src.types import mat43
 from mujoco_warp._src.types import vec5
+from mujoco_warp._src.warp_util import cache_kernel
 from mujoco_warp._src.warp_util import event_scope
 
 wp.set_module_options({"enable_backward": False})
@@ -1296,6 +1297,7 @@ def box_box_wrapper(
 }
 
 
+@cache_kernel
 def _primitive_narrowphase(primitive_collisions_types, primitive_collisions_func):
   @wp.kernel(module="unique", enable_backward=False)
   def primitive_narrowphase(

mujoco_warp/_src/constraint.py

@@ -23,6 +23,7 @@
 from mujoco_warp._src.types import DisableBit
 from mujoco_warp._src.types import vec5
 from mujoco_warp._src.types import vec11
+from mujoco_warp._src.warp_util import cache_kernel
 from mujoco_warp._src.warp_util import event_scope
 
 wp.set_module_options({"enable_backward": False})
@@ -672,6 +673,7 @@ def _equality_tendon(
   )
 
 
+@cache_kernel
 def _equality_flex(is_sparse: bool):
   @wp.kernel(module="unique", enable_backward=False)
   def kernel(

mujoco_warp/_src/forward.py

@@ -42,6 +42,7 @@
 from mujoco_warp._src.types import TileSet
 from mujoco_warp._src.types import TrnType
 from mujoco_warp._src.types import vec10f
+from mujoco_warp._src.warp_util import cache_kernel
 from mujoco_warp._src.warp_util import event_scope
 
 wp.set_module_options({"enable_backward": False})
@@ -358,6 +359,7 @@ def _euler_damp_qfrc_sparse(
   qM_integration_out[worldid, 0, adr] += timestep * damp_deriv[worldid, tid]
 
 
+@cache_kernel
 def _tile_euler_dense(tile: TileSet):
   @wp.kernel(module="unique", enable_backward=False)
   def euler_dense(

mujoco_warp/_src/sensor.py

@@ -45,6 +45,7 @@
 from mujoco_warp._src.types import vec_pluginattr
 from mujoco_warp._src.util_misc import inside_geom
 from mujoco_warp._src.util_misc import poly_potential
+from mujoco_warp._src.warp_util import cache_kernel
 from mujoco_warp._src.warp_util import event_scope
 
 wp.set_module_options({"enable_backward": False})
@@ -2406,6 +2407,7 @@ def _contact_match(
   sensor_contact_direction_out[worldid, contactsensorid, contactmatchid] = dir
 
 
+@cache_kernel
 def _contact_sort(maxmatch: int):
   @wp.kernel(module="unique", enable_backward=False)
   def contact_sort(
@@ -2901,6 +2903,7 @@ def energy_pos(m: Model, d: Data):
     # TODO(team): flex
 
 
+@cache_kernel
 def _energy_vel_kinetic(nv: int):
   @wp.kernel(module="unique", enable_backward=False)
   def energy_vel_kinetic(

mujoco_warp/_src/smooth.py

@@ -35,6 +35,7 @@
 from mujoco_warp._src.types import vec5
 from mujoco_warp._src.types import vec10
 from mujoco_warp._src.types import vec11
+from mujoco_warp._src.warp_util import cache_kernel
 from mujoco_warp._src.warp_util import event_scope
 
 wp.set_module_options({"enable_backward": False})
@@ -1063,6 +1064,7 @@ def _factor_i_sparse(m: Model, d: Data, M: wp.array3d[float], L: wp.array3d[floa
   wp.launch(_qLDiag_div, dim=(d.nworld, m.nv), inputs=[m.M_rownnz, m.M_rowadr, L], outputs=[D])
 
 
+@cache_kernel
 def _tile_cholesky_factorize(tile: TileSet):
   """Returns a kernel for dense Cholesky factorization of a tile."""
 
@@ -2694,6 +2696,7 @@ def transmission(m: Model, d: Data):
     )
 
 
+@cache_kernel
 def _solve_LD_sparse_fused(nv: int, nlevels: int):
   """Fused sparse backsubstitution: UP + diag + DOWN in one kernel."""
 
@@ -2779,6 +2782,7 @@ def _solve_LD_sparse(
   )
 
 
+@cache_kernel
 def _tile_cholesky_solve(tile: TileSet):
   """Returns a kernel for dense Cholesky backsubstitution of a tile."""
 
@@ -2856,6 +2860,7 @@ def solve_m(m: Model, d: Data, x: wp.array2d[float], y: wp.array2d[float]):
   solve_LD(m, d, d.qLD, d.qLDiagInv, x, y)
 
 
+@cache_kernel
 def _tile_cholesky_factorize_solve(tile: TileSet):
   """Returns a kernel for dense Cholesky factorization and backsubstitution of a tile."""
 

mujoco_warp/_src/solver.py

@@ -25,6 +25,7 @@
 from mujoco_warp._src import types
 from mujoco_warp._src.block_cholesky import create_blocked_cholesky_func
 from mujoco_warp._src.block_cholesky import create_blocked_cholesky_solve_func
+from mujoco_warp._src.warp_util import cache_kernel
 from mujoco_warp._src.warp_util import event_scope
 from mujoco_warp._src.warp_util import scoped_mathdx_gemm_disabled
 
@@ -882,6 +883,7 @@ def _compute_efc_eval_pt_3alphas_elliptic(
 # =============================================================================
 
 
+@cache_kernel
 def linesearch_iterative(ls_iterations: int, cone_type: types.ConeType, fuse_jv: bool, is_sparse: bool):
   """Factory for iterative linesearch kernel.
 
@@ -1408,6 +1410,7 @@ def linesearch_zero_jv(
   ctx_jv_out[worldid, efcid] = 0.0
 
 
+@cache_kernel
 def linesearch_jv_fused(is_sparse: bool, nv: int, dofs_per_thread: int):
   @wp.kernel(module="unique", enable_backward=False)
   def kernel(
@@ -1468,6 +1471,7 @@ def kernel(
   return kernel
 
 
+@cache_kernel
 def linesearch_prepare_gauss(nv: int, dofs_per_thread: int):
   @wp.kernel(module="unique", enable_backward=False)
   def kernel(
@@ -1715,6 +1719,7 @@ def solve_init_efc(
   ctx_search_dot_out[worldid] = 0.0
 
 
+@cache_kernel
 def solve_init_jaref(is_sparse: bool, nv: int, dofs_per_thread: int):
   @wp.kernel(module="unique", enable_backward=False)
   def kernel(
@@ -1797,6 +1802,7 @@ def update_constraint_init_cost(
   ctx_cost_out[worldid] = 0.0
 
 
+@cache_kernel
 def update_constraint_efc(track_changes: bool):
   TRACK_CHANGES = track_changes
 
@@ -2004,6 +2010,7 @@ def update_constraint_init_qfrc_constraint_dense(
   qfrc_constraint_out[worldid, dofid] = sum_qfrc
 
 
+@cache_kernel
 def update_constraint_gauss_cost(nv: int, dofs_per_thread: int):
   @wp.kernel(module="unique", enable_backward=False)
   def kernel(
@@ -2287,6 +2294,7 @@ def active_check(tid: int, threshold: int) -> float:
     return 1.0
 
 
+@cache_kernel
 def update_gradient_JTDAJ_sparse_tiled(tile_size: int, njmax: int):
   TILE_SIZE = tile_size
 
@@ -2356,6 +2364,7 @@ def kernel(
   return kernel
 
 
+@cache_kernel
 def update_gradient_JTDAJ_dense_tiled(nv_pad: int, tile_size: int, njmax: int):
   if njmax < tile_size:
     tile_size = njmax
@@ -2719,6 +2728,7 @@ def update_gradient_JTCJ_dense(
     ctx_h_out[worldid, dof1id, dof2id] += h
 
 
+@cache_kernel
 def update_gradient_cholesky(tile_size: int):
   @wp.kernel(module="unique", enable_backward=False)
   def kernel(
@@ -2744,6 +2754,7 @@ def kernel(
   return kernel
 
 
+@cache_kernel
 def update_gradient_cholesky_blocked(tile_size: int, matrix_size: int):
   @wp.kernel(module="unique", enable_backward=False)
   def kernel(

mujoco_warp/_src/support.py

@@ -27,6 +27,7 @@
 from mujoco_warp._src.types import State
 from mujoco_warp._src.types import vec5
 from mujoco_warp._src.types import vec10f
+from mujoco_warp._src.warp_util import cache_kernel
 from mujoco_warp._src.warp_util import event_scope
 
 wp.set_module_options({"enable_backward": False})
@@ -63,6 +64,7 @@ def next_act(
   return act
 
 
+@cache_kernel
 def mul_m_sparse(check_skip: bool):
   @wp.kernel(module="unique")
   def _mul_m_sparse(
@@ -99,6 +101,7 @@ def _mul_m_sparse(
   return _mul_m_sparse
 
 
+@cache_kernel
 def mul_m_dense(nv: int, check_skip: bool):
   """Simple SIMT dense matmul: one thread per output element."""
 
@@ -429,6 +432,7 @@ def jac_dof(
   return jacp, jacr
 
 
+@cache_kernel
 def _make_jac_kernel(has_jacp: bool, has_jacr: bool):
   @wp.kernel(module="unique", enable_backward=False)
   def _jac(

mujoco_warp/_src/warp_util.py

@@ -119,6 +119,28 @@ def wrapper(*args, **kwargs):
   return wrapper
 
 
+_KERNEL_CACHE = {}
+
+
+def cache_kernel(func):
+  # caching kernels to avoid crashes in graph_conditional code
+  @functools.wraps(func)
+  def wrapper(*args):
+    def _hash_arg(a):
+      if hasattr(a, "size"):
+        return a.size
+      if isinstance(a, list):
+        return hash(tuple(a))
+      return hash(a)
+
+    key = tuple(_hash_arg(a) for a in args) + (hash(func.__name__),)
+    if key not in _KERNEL_CACHE:
+      _KERNEL_CACHE[key] = func(*args)
+    return _KERNEL_CACHE[key]
+
+  return wrapper
+
+
 def check_toolkit_driver():
   wp.init()
   if wp.get_device().is_cuda: