Optimization: solve_LD_sparse (google-deepmind#1260)

* Fused solve_LD_sparse kernels

* Fix bug in new fused kernel

* Fix formatting

* Fix new kernel on CPU

* Change to use wp.block_dim

Author: Kenny-Vilella

mujoco_warp/_src/io.py

@@ -583,6 +583,15 @@ def _check_margin(name, t1, t2, margin):
       Madr_ki -= 1
   m.qLD_updates = tuple(wp.array(qLD_updates[i], dtype=wp.vec3i) for i in sorted(qLD_updates))
 
+  # Build concatenated updates for fused kernel
+  all_updates_flat = []
+  level_offsets = [0]
+  for level in sorted(qLD_updates):
+    all_updates_flat.extend(qLD_updates[level])
+    level_offsets.append(len(all_updates_flat))
+  m.qLD_all_updates = all_updates_flat if all_updates_flat else [(0, 0, 0)]
+  m.qLD_level_offsets = level_offsets
+
   # indices for sparse qM_fullm (used in solver)
   m.qM_fullm_i, m.qM_fullm_j = [], []
   for i in range(mjm.nv):

mujoco_warp/_src/smooth.py

@@ -2693,43 +2693,65 @@ def transmission(m: Model, d: Data):
     )
 
 
-@wp.kernel
-def _solve_LD_sparse_x_acc_up(
-  # In:
-  L: wp.array3d(dtype=float),
-  qLD_updates_: wp.array(dtype=wp.vec3i),
-  # Out:
-  x: wp.array2d(dtype=float),
-):
-  worldid, nodeid = wp.tid()
-  update = qLD_updates_[nodeid]
-  i, k, Madr_ki = update[0], update[1], update[2]
-  wp.atomic_sub(x[worldid], i, L[worldid, 0, Madr_ki] * x[worldid, k])
-
-
-@wp.kernel
-def _solve_LD_sparse_qLDiag_mul(
-  # In:
-  D: wp.array2d(dtype=float),
-  # Out:
-  out: wp.array2d(dtype=float),
-):
-  worldid, dofid = wp.tid()
-  out[worldid, dofid] *= D[worldid, dofid]
+@cache_kernel
+def _solve_LD_sparse_fused(nv: int, nlevels: int):
+  """Fused sparse backsubstitution: UP + diag + DOWN in one kernel."""
 
+  @wp.func_native(snippet="WP_TILE_SYNC();")
+  def _syncthreads():
+    pass
 
-@wp.kernel
-def _solve_LD_sparse_x_acc_down(
-  # In:
-  L: wp.array3d(dtype=float),
-  qLD_updates_: wp.array(dtype=wp.vec3i),
-  # Out:
-  x: wp.array2d(dtype=float),
-):
-  worldid, nodeid = wp.tid()
-  update = qLD_updates_[nodeid]
-  i, k, Madr_ki = update[0], update[1], update[2]
-  wp.atomic_sub(x[worldid], k, L[worldid, 0, Madr_ki] * x[worldid, i])
+  @wp.kernel(module="unique", enable_backward=False)
+  def kernel(
+    # In:
+    L: wp.array3d(dtype=float),
+    D: wp.array2d(dtype=float),
+    all_updates: wp.array(dtype=wp.vec3i),
+    level_offsets: wp.array(dtype=int),
+    y: wp.array2d(dtype=float),
+    # Out:
+    x_out: wp.array2d(dtype=float),
+  ):
+    worldid, tid = wp.tid()
+    NV = wp.static(nv)
+    NLEVELS = wp.static(nlevels)
+    BLOCK_DIM = wp.block_dim()
+
+    # Copy y to x_out
+    for dofid in range(tid, NV, BLOCK_DIM):
+      x_out[worldid, dofid] = y[worldid, dofid]
+    _syncthreads()
+
+    # Forward substitution
+    for level in range(NLEVELS):
+      level_idx = NLEVELS - 1 - level
+      level_offset = level_offsets[level_idx]
+      level_size = level_offsets[level_idx + 1] - level_offset
+
+      for u in range(tid, level_size, BLOCK_DIM):
+        update = all_updates[level_offset + u]
+        i, k, Madr_ki = update[0], update[1], update[2]
+        wp.atomic_sub(x_out[worldid], i, L[worldid, 0, Madr_ki] * x_out[worldid, k])
+      _syncthreads()
+
+    # Diagonal multiply
+    for dofid in range(tid, NV, BLOCK_DIM):
+      x_out[worldid, dofid] *= D[worldid, dofid]
+    _syncthreads()
+
+    # Backward substitution
+    for level in range(NLEVELS):
+      level_idx = level
+      level_offset = level_offsets[level_idx]
+      level_size = level_offsets[level_idx + 1] - level_offset
+
+      for u in range(tid, level_size, BLOCK_DIM):
+        update = all_updates[level_offset + u]
+        i, k, Madr_ki = update[0], update[1], update[2]
+        wp.atomic_sub(x_out[worldid], k, L[worldid, 0, Madr_ki] * x_out[worldid, i])
+      _syncthreads()
+
+  return kernel
 
 
 def _solve_LD_sparse(
@@ -2741,14 +2763,20 @@ def _solve_LD_sparse(
   y: wp.array2d(dtype=float),
 ):
   """Computes sparse backsubstitution: x = inv(L'*D*L)*y."""
-  wp.copy(x, y)
-  for qLD_updates in reversed(m.qLD_updates):
-    wp.launch(_solve_LD_sparse_x_acc_up, dim=(d.nworld, qLD_updates.size), inputs=[L, qLD_updates], outputs=[x])
-
-  wp.launch(_solve_LD_sparse_qLDiag_mul, dim=(d.nworld, m.nv), inputs=[D], outputs=[x])
+  nlevels = len(m.qLD_updates)
+  if wp.get_device().is_cuda:
+    dim_block = m.block_dim.solve_LD_sparse_fused
+  else:
+    # Fallback for CPU
+    dim_block = 1
 
-  for qLD_updates in m.qLD_updates:
-    wp.launch(_solve_LD_sparse_x_acc_down, dim=(d.nworld, qLD_updates.size), inputs=[L, qLD_updates], outputs=[x])
+  wp.launch(
+    _solve_LD_sparse_fused(m.nv, nlevels),
+    dim=(d.nworld, dim_block),
+    inputs=[L, D, m.qLD_all_updates, m.qLD_level_offsets, y],
+    outputs=[x],
+    block_dim=dim_block,
+  )
 
 
 @cache_kernel

mujoco_warp/_src/types.py

@@ -59,6 +59,7 @@ class BlockDim:
   cholesky_factorize: int = 32
   cholesky_solve: int = 32
   cholesky_factorize_solve: int = 32
+  solve_LD_sparse_fused: int = 64
   # solver
   update_gradient_cholesky: int = 64
   update_gradient_cholesky_blocked: int = 32
@@ -1204,6 +1205,8 @@ class Model:
     taxel_sensorid: address for tactile sensors
     qM_tiles: tiling configuration
     qLD_updates: tuple of index triples for sparse factorization
+    qLD_all_updates: tuple of all levels concatenated
+    qLD_level_offsets: tuple of start offsets for each level
     qM_fullm_i: sparse mass matrix addressing
     qM_fullm_j: sparse mass matrix addressing
     qM_mulm_rowadr: sparse matmul row pointers
@@ -1580,6 +1583,8 @@ class Model:
   taxel_sensorid: wp.array(dtype=int)
   qM_tiles: tuple[TileSet, ...]
   qLD_updates: tuple[wp.array(dtype=wp.vec3i), ...]
+  qLD_all_updates: wp.array(dtype=wp.vec3i)
+  qLD_level_offsets: wp.array(dtype=int)
   qM_fullm_i: wp.array(dtype=int)
   qM_fullm_j: wp.array(dtype=int)
   # Gather-based sparse mul_m indices (thread per DOF, no atomics)