chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge branch 'google-deepmind:main' into mjx-warp-segmentation

Browse Source
This commit is contained in:
Tarik Kelestemur
2026-04-26 21:36:32 -04:00
committed by GitHub
parent 22f4458b1c b643f50b2a
commit 1009478de9
75 changed files with 2344 additions and 1124 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/io.py
+13 -17
View File
@@ -96,16 +96,15 @@ def _resolve_device(
return cpu_0
if impl == types.Impl.WARP:
# WARP implementation requires a CUDA GPU.
cuda_gpus = [d for d in jax.devices('cuda')]
if not cuda_gpus:
raise AssertionError(
'No CUDA GPU devices found in'
f' jax.devices("cuda")={jax.devices("cuda")}.'
)
# WARP implementation requires a CUDA GPU or CPU.
if has_cuda_gpu_device():
cuda_gpus = [d for d in jax.devices('cuda')]
if cuda_gpus:
logging.debug('Picking default device: %s', cuda_gpus[0])
return cuda_gpus[0]
logging.debug('Picking default device: %s', cuda_gpus[0])
return cuda_gpus[0]
logging.debug('Picking default device for Warp: CPU')
return jax.devices('cpu')[0]
raise ValueError(f'Unsupported implementation: {impl}')
@@ -121,9 +120,12 @@ def _check_impl_device_compatibility(
impl = types.Impl(impl)
if impl == types.Impl.WARP:
if not _is_cuda_gpu_device(device):
is_cuda_device = _is_cuda_gpu_device(device)
is_cpu_device = device.platform == 'cpu'
if not (is_cuda_device or is_cpu_device):
raise AssertionError(
f'Warp implementation requires a CUDA GPU device, got {device}.'
'Warp implementation requires a CUDA GPU or CPU device, got '
f'{device}.'
)
_check_warp_installed()
@@ -425,8 +427,6 @@ def _put_model_jax(
return _strip_weak_type(model)
def _put_model_warp(
m: mujoco.MjModel,
graph_mode: mjxw.types.GraphMode,
@@ -719,8 +719,6 @@ def _make_data_jax(
return d
def _get_nested_attr(obj: Any, attr_name: str, split: str) -> Any:
"""Returns the nested attribute from an object."""
for part in attr_name.split(split):
@@ -1084,8 +1082,6 @@ def _put_data_jax(
return _strip_weak_type(data)
# TODO(josechenf): Iterate on the keepalive implementation to make it easier to
# use before OSS.
def _put_data_cpp(
mjx/mujoco/mjx/_src/io_test.py
+5 -14
View File
@@ -934,8 +934,8 @@ _DEVICE_TEST_CASES = [
('gpu-nvidia', 'jax', ('gpu', Impl.JAX)),
('tpu', 'jax', ('tpu', Impl.JAX)),
# WARP backend specified.
('cpu', 'warp', ('cpu', 'error')),
('gpu-notnvidia', 'warp', ('cpu', 'error')),
('cpu', 'warp', ('cpu', Impl.WARP)),
('gpu-notnvidia', 'warp', ('gpu', 'error')),
('gpu-nvidia', 'warp', ('gpu', Impl.WARP)),
('tpu', 'warp', ('tpu', 'error')),
# CPP backend specified.
@@ -962,10 +962,10 @@ _DEFAULT_DEVICE_TEST_CASES = [
('gpu-nvidia', 'jax', ('gpu', Impl.JAX)),
('tpu', 'jax', ('tpu', Impl.JAX)),
# WARP backend impl specified.
('cpu', 'warp', ('cpu', 'error')),
('gpu-notnvidia', 'warp', ('cpu', 'error')),
('cpu', 'warp', ('cpu', Impl.WARP)),
('gpu-notnvidia', 'warp', ('cpu', Impl.WARP)),
('gpu-nvidia', 'warp', ('gpu', Impl.WARP)),
('tpu', 'warp', ('tpu', 'error')),
('tpu', 'warp', ('cpu', Impl.WARP)),
# CPP backend impl specified, CPU should always be available.
('cpu', 'cpp', ('cpu', Impl.CPP)),
('gpu-notnvidia', 'cpp', ('cpu', Impl.CPP)),
@@ -1140,15 +1140,6 @@ class ResolveImplAndDeviceTest(parameterized.TestCase):
self.mock_jax_backends.side_effect = backends_side_effect
expected_device, expected_impl = expected
if (
expected_impl == 'error'
and default_device_str != 'gpu-nvidia'
and impl_str == 'warp'
):
with self.assertRaisesRegex(RuntimeError, 'cuda backend not supported'):
mjx_io._resolve_impl_and_device(impl=impl_str, device=None)
return
if expected_impl == 'error':
with self.assertRaises(AssertionError):
mjx_io._resolve_impl_and_device(impl=impl_str, device=None)
mjx/mujoco/mjx/_src/types.py
+4
View File
@@ -776,6 +776,10 @@ class Model(PyTreeNode):
flex_vertadr: np.ndarray
flex_vertnum: np.ndarray
flex_vert0: np.ndarray
flex_nodeadr: np.ndarray
flex_nodenum: np.ndarray
flex_nodebodyid: np.ndarray
flex_node0: np.ndarray
hfield_size: np.ndarray
hfield_nrow: np.ndarray
hfield_ncol: np.ndarray
mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/ffi.py
Vendored
+111 -48
View File
@@ -210,11 +210,19 @@ class FfiKernel:
self.input_output_aliases = input_output_aliases
# register the callback
FFI_CCALLFUNC = ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.POINTER(XLA_FFI_CallFrame))
self.callback_func = FFI_CCALLFUNC(self.ffi_callback)
ffi_ccall_address = ctypes.cast(self.callback_func, ctypes.c_void_p)
ffi_capsule = jax.ffi.pycapsule(ffi_ccall_address.value)
jax.ffi.register_ffi_target(self.name, ffi_capsule, platform="CUDA")
FFI_CCALLFUNC = ctypes.CFUNCTYPE(
ctypes.c_void_p, ctypes.POINTER(XLA_FFI_CallFrame)
)
self.callback_func_cuda = FFI_CCALLFUNC(lambda call_frame: self.ffi_callback(call_frame, platform="CUDA"))
ffi_ccall_address_cuda = ctypes.cast(self.callback_func_cuda, ctypes.c_void_p)
ffi_capsule_cuda = jax.ffi.pycapsule(ffi_ccall_address_cuda.value)
jax.ffi.register_ffi_target(self.name, ffi_capsule_cuda, platform="CUDA")
self.callback_func_host = FFI_CCALLFUNC(lambda call_frame: self.ffi_callback(call_frame, platform="Host"))
ffi_ccall_address_host = ctypes.cast(self.callback_func_host, ctypes.c_void_p)
ffi_capsule_host = jax.ffi.pycapsule(ffi_ccall_address_host.value)
jax.ffi.register_ffi_target(self.name, ffi_capsule_host, platform="Host")
def __call__(self, *args, output_dims=None, launch_dims=None, vmap_method=None):
num_inputs = len(args)
@@ -241,18 +249,19 @@ class FfiKernel:
# check dtype
if input_value.dtype != input_arg.jax_scalar_type:
raise TypeError(
f"Invalid data type for array argument '{input_arg.name}', expected {input_arg.jax_scalar_type}, got {input_value.dtype}"
f"Invalid data type for array argument '{input_arg.name}',"
f" expected {input_arg.jax_scalar_type}, got {input_value.dtype}"
)
# check ndim
if input_value.ndim != input_arg.jax_ndim:
raise TypeError(
f"Invalid dimensionality for array argument '{input_arg.name}', expected {input_arg.jax_ndim} dimensions, got {input_value.ndim}"
f"Invalid dimensionality for array argument '{input_arg.name}', expected {input_arg.jax_ndim} dimensions, got {input_value.ndim}"
)
# check inner dims
for d in range(input_arg.dtype_ndim):
if input_value.shape[input_arg.type.ndim + d] != input_arg.dtype_shape[d]:
raise TypeError(
f"Invalid inner dimensions for array argument '{input_arg.name}', expected {input_arg.dtype_shape}, got {input_value.shape[-input_arg.dtype_ndim :]}"
f"Invalid inner dimensions for array argument '{input_arg.name}', expected {input_arg.dtype_shape}, got {input_value.shape[-input_arg.dtype_ndim :]}"
)
else:
# make sure scalar is not a traced variable, should be static
@@ -328,7 +337,7 @@ class FfiKernel:
return call(*args, launch_id=launch_id)
def ffi_callback(self, call_frame):
def ffi_callback(self, call_frame, platform="CUDA"):
try:
# On the first call, XLA runtime will query the API version and traits
# metadata using the |extension| field. Let us respond to that query
@@ -340,10 +349,11 @@ class FfiKernel:
metadata_ext = ctypes.cast(extension, ctypes.POINTER(XLA_FFI_Metadata_Extension))
metadata_ext.contents.metadata.contents.api_version.major_version = 0
metadata_ext.contents.metadata.contents.api_version.minor_version = 1
# Turn on CUDA graphs for this handler.
metadata_ext.contents.metadata.contents.traits = (
XLA_FFI_Handler_TraitsBits.COMMAND_BUFFER_COMPATIBLE
)
# Turn on CUDA graphs for this handler if on CUDA platform.
if platform == "CUDA":
metadata_ext.contents.metadata.contents.traits = (
XLA_FFI_Handler_TraitsBits.COMMAND_BUFFER_COMPATIBLE
)
return None
# Lock is required to prevent race conditions when callback is invoked
@@ -423,29 +433,43 @@ class FfiKernel:
kernel_params[0] = ctypes.addressof(launch_bounds)
# get device and stream
device = wp.get_cuda_device(get_device_ordinal_from_callframe(call_frame.contents))
stream = get_stream_from_callframe(call_frame.contents)
if platform == "CUDA":
device = wp.get_cuda_device(get_device_ordinal_from_callframe(call_frame.contents))
stream = get_stream_from_callframe(call_frame.contents)
else:
device = wp.get_device("cpu")
stream = None
# get kernel hooks
hooks = self.kernel.module.get_kernel_hooks(self.kernel, device)
assert hooks.forward, "Failed to find kernel entry point"
# launch the kernel
wp._src.context.runtime.core.wp_cuda_launch_kernel(
device.context,
hooks.forward,
launch_bounds.size,
0,
256,
hooks.forward_smem_bytes,
kernel_params,
stream,
)
if device.is_cuda:
wp._src.context.runtime.core.wp_cuda_launch_kernel(
device.context,
hooks.forward,
launch_bounds.size,
0,
256,
hooks.forward_smem_bytes,
kernel_params,
stream,
)
else:
wp._src.context.runtime.core.wp_cpu_launch_kernel(
device.context,
hooks.forward,
launch_bounds.size,
kernel_params,
)
except Exception as e:
print(traceback.format_exc())
return create_ffi_error(
call_frame.contents.api, XLA_FFI_Error_Code.UNKNOWN, f"FFI callback error: {type(e).__name__}: {e}"
call_frame.contents.api,
XLA_FFI_Error_Code.UNKNOWN,
f"FFI callback error: {type(e).__name__}: {e}",
)
@@ -594,10 +618,16 @@ class FfiCallable:
# register the callback
FFI_CCALLFUNC = ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.POINTER(XLA_FFI_CallFrame))
self.callback_func = FFI_CCALLFUNC(self.ffi_callback)
ffi_ccall_address = ctypes.cast(self.callback_func, ctypes.c_void_p)
ffi_capsule = jax.ffi.pycapsule(ffi_ccall_address.value)
jax.ffi.register_ffi_target(self.name, ffi_capsule, platform="CUDA")
self.callback_func_cuda = FFI_CCALLFUNC(lambda call_frame: self.ffi_callback(call_frame, platform="CUDA"))
ffi_ccall_address_cuda = ctypes.cast(self.callback_func_cuda, ctypes.c_void_p)
ffi_capsule_cuda = jax.ffi.pycapsule(ffi_ccall_address_cuda.value)
jax.ffi.register_ffi_target(self.name, ffi_capsule_cuda, platform="CUDA")
self.callback_func_host = FFI_CCALLFUNC(lambda call_frame: self.ffi_callback(call_frame, platform="Host"))
ffi_ccall_address_host = ctypes.cast(self.callback_func_host, ctypes.c_void_p)
ffi_capsule_host = jax.ffi.pycapsule(ffi_ccall_address_host.value)
jax.ffi.register_ffi_target(self.name, ffi_capsule_host, platform="Host")
def __call__(self, *args, output_dims=None, vmap_method=None):
num_inputs = len(args)
@@ -688,8 +718,7 @@ class FfiCallable:
except Exception:
# ignore unsupported devices like TPUs
pass
# we only support CUDA devices for now
if dev.is_cuda:
if dev.is_cuda or dev.is_cpu:
module.load(dev)
# save call data to be retrieved by callback
@@ -698,7 +727,7 @@ class FfiCallable:
self.call_id += 1
return call(*args, call_id=call_id)
def ffi_callback(self, call_frame):
def ffi_callback(self, call_frame, platform="CUDA"):
try:
# On the first call, XLA runtime will query the API version and traits
# metadata using the |extension| field. Let us respond to that query
@@ -710,8 +739,8 @@ class FfiCallable:
metadata_ext = ctypes.cast(extension, ctypes.POINTER(XLA_FFI_Metadata_Extension))
metadata_ext.contents.metadata.contents.api_version.major_version = 0
metadata_ext.contents.metadata.contents.api_version.minor_version = 1
# Turn on CUDA graphs for this handler.
if self.graph_mode is GraphMode.JAX:
# Turn on CUDA graphs for this handler if on CUDA platform.
if self.graph_mode is GraphMode.JAX and platform == "CUDA":
metadata_ext.contents.metadata.contents.traits = (
XLA_FFI_Handler_TraitsBits.COMMAND_BUFFER_COMPATIBLE
)
@@ -738,6 +767,35 @@ class FfiCallable:
assert num_inputs == self.num_inputs
assert num_outputs == self.num_outputs
if platform == "Host":
device = wp.get_device("cpu")
# reconstruct the argument list
arg_list = []
# input and in-out args
for i, arg in enumerate(self.input_args):
if arg.is_array:
buffer = inputs[i].contents
shape = collapse_batch_dims(buffer.dims[: buffer.rank - arg.dtype_ndim], arg.type.ndim)
arr = wp.array(ptr=buffer.data, dtype=arg.type.dtype, shape=shape, device=device)
arg_list.append(arr)
else:
# scalar argument, get stashed value
value = call_desc.static_inputs[arg.name]
arg_list.append(value)
# pure output args (skip in-out FFI buffers)
for i, arg in enumerate(self.output_args):
buffer = outputs[i + self.num_in_out].contents
shape = collapse_batch_dims(buffer.dims[: buffer.rank - arg.dtype_ndim], arg.type.ndim)
arr = wp.array(ptr=buffer.data, dtype=arg.type.dtype, shape=shape, device=device)
arg_list.append(arr)
# call the Python function with reconstructed arguments
with wp.ScopedDevice(device):
self.func(*arg_list)
return
cuda_stream = get_stream_from_callframe(call_frame.contents)
device_ordinal = get_device_ordinal_from_callframe(call_frame.contents)
@@ -870,8 +928,8 @@ class FfiCallable:
arg_list.append(arr)
# call the Python function with reconstructed arguments
with wp.ScopedStream(stream, sync_enter=False):
if stream.is_capturing:
with wp.ScopedStream(stream, sync_enter=False) if stream else wp.ScopedDevice(device):
if stream and stream.is_capturing:
# capturing with JAX
with wp.ScopedCapture(external=True) as capture:
self.func(*arg_list)
@@ -879,7 +937,7 @@ class FfiCallable:
# keep a reference to the capture object to prevent required modules getting unloaded
call_desc.capture = capture
elif self.graph_mode == GraphMode.WARP:
elif self.graph_mode == GraphMode.WARP and device.is_cuda:
# capturing with WARP
with wp.ScopedCapture() as capture:
self.func(*arg_list)
@@ -892,7 +950,7 @@ class FfiCallable:
if self._graph_cache_max is not None and len(self.captures) > self._graph_cache_max:
self.captures.popitem(last=False)
elif self.graph_mode == GraphMode.WARP_STAGED_EX:
elif self.graph_mode == GraphMode.WARP_STAGED_EX and device.is_cuda:
# capturing with WARP using staging buffers and memcopies done outside of the graph
wp_memcpy_batch = wp._src.context.runtime.core.wp_memcpy_batch
@@ -935,7 +993,7 @@ class FfiCallable:
# TODO: we should have a way of freeing this
call_desc.capture = capture
elif self.graph_mode == GraphMode.WARP_STAGED:
elif self.graph_mode == GraphMode.WARP_STAGED and device.is_cuda:
# capturing with WARP using staging buffers and memcopies done inside of the graph
wp_cuda_graph_insert_memcpy_batch = (
wp._src.context.runtime.core.wp_cuda_graph_insert_memcpy_batch
@@ -1013,7 +1071,7 @@ class FfiCallable:
call_desc.capture = capture
else:
# not capturing
# not capturing or on CPU
self.func(*arg_list)
except Exception as e:
@@ -1621,7 +1679,7 @@ def register_ffi_callback(name: str, func: Callable, graph_compatible: bool = Tr
# TODO check that the name is not already registered
def ffi_callback(call_frame):
def ffi_callback(call_frame, platform="CUDA"):
try:
extension = call_frame.contents.extension_start
# On the first call, XLA runtime will query the API version and traits
@@ -1633,7 +1691,7 @@ def register_ffi_callback(name: str, func: Callable, graph_compatible: bool = Tr
metadata_ext = ctypes.cast(extension, ctypes.POINTER(XLA_FFI_Metadata_Extension))
metadata_ext.contents.metadata.contents.api_version.major_version = 0
metadata_ext.contents.metadata.contents.api_version.minor_version = 1
if graph_compatible:
if graph_compatible and platform == "CUDA":
# Turn on CUDA graphs for this handler.
metadata_ext.contents.metadata.contents.traits = (
XLA_FFI_Handler_TraitsBits.COMMAND_BUFFER_COMPATIBLE
@@ -1666,12 +1724,17 @@ def register_ffi_callback(name: str, func: Callable, graph_compatible: bool = Tr
return None
FFI_CCALLFUNC = ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.POINTER(XLA_FFI_CallFrame))
callback_func = FFI_CCALLFUNC(ffi_callback)
callback_func_cuda = FFI_CCALLFUNC(lambda call_frame: ffi_callback(call_frame, platform="CUDA"))
callback_func_host = FFI_CCALLFUNC(lambda call_frame: ffi_callback(call_frame, platform="Host"))
with _FFI_REGISTRY_LOCK:
_FFI_CALLBACK_REGISTRY[name] = callback_func
ffi_ccall_address = ctypes.cast(callback_func, ctypes.c_void_p)
ffi_capsule = jax.ffi.pycapsule(ffi_ccall_address.value)
jax.ffi.register_ffi_target(name, ffi_capsule, platform="CUDA")
_FFI_CALLBACK_REGISTRY[f"{name}_cuda"] = callback_func_cuda
_FFI_CALLBACK_REGISTRY[f"{name}_host"] = callback_func_host
ffi_ccall_address_cuda = ctypes.cast(callback_func_cuda, ctypes.c_void_p)
ffi_capsule_cuda = jax.ffi.pycapsule(ffi_ccall_address_cuda.value)
jax.ffi.register_ffi_target(name, ffi_capsule_cuda, platform="CUDA")
ffi_ccall_address_host = ctypes.cast(callback_func_host, ctypes.c_void_p)
ffi_capsule_host = jax.ffi.pycapsule(ffi_ccall_address_host.value)
jax.ffi.register_ffi_target(name, ffi_capsule_host, platform="Host")
###############################################################################
mjx/mujoco/mjx/warp/forward_test.py
+51
View File
@@ -31,6 +31,7 @@ from mujoco.mjx.warp import test_util as tu
from mujoco.mjx.warp import warp as wp # pylint: disable=g-importing-member
import numpy as np
try:
from mujoco.mjx.warp import forward # pylint: disable=g-import-not-at-top
except ImportError:
@@ -300,6 +301,56 @@ class StepTest(parameterized.TestCase):
tu.assert_attr_eq(dx, d, 'mocap_quat')
tu.assert_attr_eq(dx, d, 'sensordata')
@parameterized.parameters(
'humanoid/humanoid.xml',
'pendula.xml',
)
def test_step_cpu(self, xml: str):
"""Tests step on the CPU device."""
if not _FORCE_TEST:
if not mjxw.WARP_INSTALLED:
self.skipTest('Warp not installed.')
batch_size = 1
m = test_util.load_test_file(xml)
m.opt.iterations = 10
m.opt.ls_iterations = 10
cpu_device = jax.devices('cpu')[0]
mx = mjx.put_model(m, impl='warp', device=cpu_device)
d = mujoco.MjData(m)
worldids = jp.arange(batch_size)
dx_batch = jax.vmap(functools.partial(tu.make_data, m))(worldids)
dx_batch = jax.device_put(dx_batch, cpu_device)
dx_batch_orig = dx_batch
for _ in range(10):
dx_batch = jax.vmap(forward.step, in_axes=(None, 0))(
mx, dx_batch
)
for i in range(batch_size):
dx = dx_batch[i]
dx_orig = dx_batch_orig[i]
d.qpos[:] = dx_orig.qpos
d.qvel[:] = dx_orig.qvel
d.ctrl[:] = dx_orig.ctrl
d.mocap_pos[:] = dx_orig.mocap_pos
d.mocap_quat[:] = dx_orig.mocap_quat
d.time = dx_orig.time
mujoco.mj_step(m, d, 10)
tu.assert_attr_eq(dx, d, 'qpos')
tu.assert_attr_eq(dx, d, 'qvel')
tu.assert_attr_eq(dx, d, 'time')
tu.assert_attr_eq(dx, d, 'ctrl')
tu.assert_attr_eq(dx, d, 'act')
tu.assert_attr_eq(dx, d, 'mocap_pos')
tu.assert_attr_eq(dx, d, 'mocap_quat')
tu.assert_attr_eq(dx, d, 'sensordata')
def test_step_leading_dim_mismatch(self):
if not _FORCE_TEST:
if not mjxw.WARP_INSTALLED:
mjx/mujoco/mjx/warp/smooth_test.py
+5 -4
View File
@@ -133,10 +133,11 @@ class SmoothTest(parameterized.TestCase):
def test_kinematics_vmap(self):
"""Tests kinematics with batched data."""
if not mjxw.WARP_INSTALLED:
self.skipTest('Warp not installed.')
if not io.has_cuda_gpu_device():
self.skipTest('No CUDA GPU device available.')
if not _FORCE_TEST:
if not mjxw.WARP_INSTALLED:
self.skipTest('Warp not installed.')
if not io.has_cuda_gpu_device():
self.skipTest('No CUDA GPU device available.')
m = tu.load_test_file('pendula.xml')
mjx/mujoco/mjx/warp/test_util.py
+6 -2
View File
@@ -153,7 +153,9 @@ def _mjx_efc(dx, worldid: int):
efc_pos = select(dx._impl.efc__pos)[:nefc]
efc_type = select(dx._impl.efc__type)[:nefc]
efc_d = select(dx._impl.efc__D)[:nefc]
keys_sorted = np.lexsort((-efc_pos, efc_type, efc_d))
keys_sorted = np.lexsort(
(-np.round(efc_pos, 12), efc_type, np.round(efc_d, 12))
)
keys = keys[keys_sorted]
nefc = len(keys)
@@ -180,7 +182,9 @@ def _mj_efc(d):
else:
efc_j = d.efc_J.reshape((-1, d.qvel.shape[0]))
keys = np.lexsort((-d.efc_pos, d.efc_type, d.efc_D))
keys = np.lexsort(
(-np.round(d.efc_pos, 12), d.efc_type, np.round(d.efc_D, 12))
)
type_ = d.efc_type[keys]
pos = d.efc_pos[keys]
efc_j = efc_j[keys]