diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index 74937bc0..918b0fa4 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -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( diff --git a/mjx/mujoco/mjx/_src/io_test.py b/mjx/mujoco/mjx/_src/io_test.py index dfcb541d..542f0718 100644 --- a/mjx/mujoco/mjx/_src/io_test.py +++ b/mjx/mujoco/mjx/_src/io_test.py @@ -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) diff --git a/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/ffi.py b/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/ffi.py index c3f9e6e9..d931dded 100644 --- a/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/ffi.py +++ b/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/ffi.py @@ -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") ############################################################################### diff --git a/mjx/mujoco/mjx/warp/forward_test.py b/mjx/mujoco/mjx/warp/forward_test.py index b80c921d..056a0c51 100644 --- a/mjx/mujoco/mjx/warp/forward_test.py +++ b/mjx/mujoco/mjx/warp/forward_test.py @@ -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: diff --git a/mjx/mujoco/mjx/warp/smooth_test.py b/mjx/mujoco/mjx/warp/smooth_test.py index e7ffb92b..7c8bb368 100644 --- a/mjx/mujoco/mjx/warp/smooth_test.py +++ b/mjx/mujoco/mjx/warp/smooth_test.py @@ -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') diff --git a/mjx/mujoco/mjx/warp/test_util.py b/mjx/mujoco/mjx/warp/test_util.py index 1752f3c9..78a28671 100644 --- a/mjx/mujoco/mjx/warp/test_util.py +++ b/mjx/mujoco/mjx/warp/test_util.py @@ -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]