Implementation:Google deepmind Mujoco MJWarp Block Cholesky
| Knowledge Sources | |
|---|---|
| Domains | Physics_Simulation, GPU_Computing, NVIDIA_Warp |
| Last Updated | 2026-02-15 04:00 GMT |
Overview
MJWarp Block Cholesky provides blocked Cholesky factorization and solve functions using Warp tile operations for dense matrix computations on the GPU.
Description
This module implements blocked Cholesky decomposition (A = L * L^T) and blocked triangular solve (L * L^T * x = b) using NVIDIA Warp's tile-based shared memory operations. The functions are generated via @lru_cache-decorated factory functions that produce specialized Warp functions for given block sizes and matrix dimensions. The blocked approach processes the matrix in tiles using wp.tile_load, wp.tile_matmul, wp.tile_cholesky, wp.tile_lower_solve_inplace, and wp.tile_upper_solve_inplace for efficient GPU execution with shared memory.
Usage
Used by the solver module for dense mass matrix factorization and constraint gradient computation when the model uses dense Jacobian representation or when blocked Cholesky is selected for the gradient update.
Code Reference
Source Location
- Repository: Google_deepmind_Mujoco
- File: mjx/mujoco/mjx/third_party/mujoco_warp/_src/block_cholesky.py
- Lines: 1-108
Key Functions
# Factory function for blocked Cholesky factorization
@lru_cache(maxsize=None)
def create_blocked_cholesky_func(block_size: int):
@wp.func
def blocked_cholesky_func(A, matrix_size, L):
"""Computes L such that A = L L^T using blocked operations."""
...
# Factory function for blocked Cholesky solve
@lru_cache(maxsize=None)
def create_blocked_cholesky_solve_func(block_size: int, matrix_size_static: int):
@wp.func
def blocked_cholesky_solve_func(L, b, matrix_size, x):
"""Solves A x = b given Cholesky factor L using forward/backward substitution."""
...
Import
from mujoco.mjx.third_party.mujoco_warp._src.block_cholesky import create_blocked_cholesky_func
from mujoco.mjx.third_party.mujoco_warp._src.block_cholesky import create_blocked_cholesky_solve_func
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| block_size | int | Yes | Block dimension for tiled operations (e.g., 32) |
| A | wp.array(dtype=float, ndim=2) | Yes | Symmetric positive definite matrix to factorize |
| matrix_size | int | Yes | Dimension of the matrix |
| b | wp.array(dtype=float, ndim=2) | No | Right-hand side vector for solve |
Outputs
| Name | Type | Description |
|---|---|---|
| L | wp.array(dtype=float, ndim=2) | Lower-triangular Cholesky factor |
| x | wp.array(dtype=float, ndim=2) | Solution vector of L L^T x = b |