From 0b54638df486242b54953274b105282b0b157f8a Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Wed, 1 Apr 2026 01:53:33 -0700 Subject: [PATCH] Import NVIDIA/warp from GitHub. PiperOrigin-RevId: 892794388 Change-Id: Ic681071dc68bfaecaaab4e131c1c464dca99f9b5 --- mjx/cuda_requirements.txt | 10 +- .../warp/_src/jax_experimental/custom_call.py | 6 +- .../warp/_src/jax_experimental/ffi.py | 178 ++++++++++++++---- mjx/pyproject.toml | 2 +- 4 files changed, 152 insertions(+), 44 deletions(-) diff --git a/mjx/cuda_requirements.txt b/mjx/cuda_requirements.txt index 0f9ae0d5..2ced5d12 100644 --- a/mjx/cuda_requirements.txt +++ b/mjx/cuda_requirements.txt @@ -16,8 +16,8 @@ jax-cuda12-pjrt==0.5.3; python_version >= '3.10' \ jax-cuda12-pjrt==0.4.30; python_version == '3.9' \ --hash=sha256:895d0198ad99638fcaf976c47592e2a543eef79ea15fabd24a402d055390c328 \ --hash=sha256:c36fb1e0c236563bf3a87e70f4d1ab28a31d7cf5d722c9ede30c4172116e8bcb -warp-lang==1.11.1 \ - --hash=sha256:1ad11f1fa775269e991a3d55039152c8a504baf86701c849b485cb8e66c49d15 \ - --hash=sha256:8b098f41e71d421d80ee7562e38aa8380ff6b0d3b4c6ee866cfbdef733ac5bdc \ - --hash=sha256:5d0904b0eefcc81f39ba65375427a3de99006088aa43e24a9011263f07d0cd07 \ - --hash=sha256:15dc10aa51fb0fdbe1ca16d52e5fadca35a47ffd9d0c636826506f96bb2e7c41 +warp-lang==1.12.0 \ + --hash=sha256:c78c3701d5cad86c30ef5017410d294ec46a396bb0d502ee1c98743494f3a62f \ + --hash=sha256:a1436f60a1881cd94f787e751a83fc0987626be2d3e2b4e74c64a6947c6d1266 \ + --hash=sha256:a2d6decba693aba5b828573c4414fd6a3f4c4a934db9c322736ef2b3fa99fe76 \ + --hash=sha256:697248edd2f1e2952f50e3db33b214af76173641a8894aacc467bed6dc247f8a diff --git a/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/custom_call.py b/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/custom_call.py index b46a1071..0adf6435 100644 --- a/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/custom_call.py +++ b/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/custom_call.py @@ -19,7 +19,7 @@ from functools import reduce import warp as wp from warp._src.context import type_str from warp._src.jax import get_jax_device -from warp._src.types import array_t, launch_bounds_t, strides_from_shape +from warp._src.types import array_t, launch_bounds_t, matches_array_class, strides_from_shape from warp._src.utils import warn _wp_module_name_ = "warp.jax_experimental.custom_call" @@ -340,7 +340,7 @@ def _create_jax_warp_primitive(): wtype = warg.type rtt = ir.RankedTensorType(actual.type) - if not isinstance(wtype, wp.array): + if not matches_array_class(wtype, wp.array): raise Exception("Only contiguous arrays are supported for Jax kernel arguments") if not base_type_is_compatible(wtype.dtype, rtt.element_type): @@ -364,7 +364,7 @@ def _create_jax_warp_primitive(): for warg in wp_kernel.adj.args[len(args) :]: wtype = warg.type - if not isinstance(wtype, wp.array): + if not matches_array_class(wtype, wp.array): raise Exception("Only contiguous arrays are supported for Jax kernel arguments") # Infer dimensions from the first input. 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 f5c925dd..e9fe408f 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 @@ -29,19 +29,23 @@ import warp as wp from warp._src.codegen import get_full_arg_spec, make_full_qualified_name from warp._src.context import CudaMemcpyKind from warp._src.jax import get_jax_device -from warp._src.types import array_t, launch_bounds_t, strides_from_shape, type_size_in_bytes, type_to_warp +from warp._src.types import ( + array_t, + launch_bounds_t, + matches_array_class, + strides_from_shape, + type_size_in_bytes, + type_to_warp, +) from .xla_ffi import * _wp_module_name_ = "warp.jax_experimental.ffi" -# Type alias for differentiable kernel cache key -DiffKernelCacheKey = tuple[Callable, tuple, int, str, tuple[str, ...]] - # Holders for the custom callbacks to keep them alive. -_FFI_KERNEL_REGISTRY: dict[str, FfiKernel] = {} -_FFI_DIFF_KERNEL_REGISTRY: dict[DiffKernelCacheKey, Callable] = {} -_FFI_CALLABLE_REGISTRY: dict[str, FfiCallable] = {} +_FFI_KERNEL_REGISTRY: dict[tuple, FfiKernel] = {} +_FFI_DIFF_KERNEL_REGISTRY: dict[tuple, Callable] = {} +_FFI_CALLABLE_REGISTRY: dict[tuple, FfiCallable] = {} _FFI_CALLBACK_REGISTRY: dict[str, ctypes.CFUNCTYPE] = {} _FFI_REGISTRY_LOCK = threading.Lock() @@ -61,6 +65,21 @@ def check_jax_version(): raise RuntimeError(msg) +def collapse_batch_dims(shape, desired_ndim): + # roll leading batch dims into one + while len(shape) > desired_ndim: + shape = (shape[0] * shape[1], *shape[2:]) + return shape + + +def compute_batch_size(shape, batch_ndim): + # compute product of batch dims at front + batch_size = 1 + for i in range(batch_ndim): + batch_size *= shape[i] + return batch_size + + class GraphMode(IntEnum): """CUDA graph capture modes for :func:`warp.jax_experimental.jax_callable`. @@ -91,7 +110,7 @@ class FfiArg: self.name = name self.type = type self.in_out = in_out - self.is_array = isinstance(type, wp.array) + self.is_array = matches_array_class(type, wp.array) if self.is_array: if hasattr(type.dtype, "_wp_scalar_type_"): @@ -125,7 +144,15 @@ class FfiLaunchDesc: class FfiKernel: def __init__( - self, kernel, num_outputs, vmap_method, launch_dims, output_dims, in_out_argnames, module_preload_mode + self, + kernel, + num_outputs, + vmap_method, + launch_dims, + output_dims, + in_out_argnames, + module_preload_mode, + has_side_effect=False, ): self.kernel = kernel self.name = generate_unique_name(kernel.func) @@ -134,6 +161,7 @@ class FfiKernel: self.launch_dims = launch_dims self.output_dims = output_dims self.module_preload_mode = module_preload_mode + self.has_side_effect = has_side_effect self.first_array_arg = None self.launch_id = 0 self.launch_descriptors = {} @@ -250,7 +278,8 @@ class FfiKernel: out_types.append(get_jax_output_type(input_arg, input_value.shape)) # launch dimensions - if launch_dims is None: + infer_launch_dims = launch_dims is None + if infer_launch_dims: # use the shape of the first input array if self.first_array_arg is not None: launch_dims = get_warp_shape(self.input_args[self.first_array_arg], args[self.first_array_arg].shape) @@ -284,6 +313,7 @@ class FfiKernel: out_types, vmap_method=vmap_method, input_output_aliases=self.input_output_aliases, + has_side_effect=self.has_side_effect, ) # preload on the specified devices @@ -303,7 +333,9 @@ class FfiKernel: # save launch data to be retrieved by callback launch_id = self.launch_id - self.launch_descriptors[launch_id] = FfiLaunchDesc(static_inputs, launch_dims) + self.launch_descriptors[launch_id] = FfiLaunchDesc( + static_inputs, launch_dims if not infer_launch_dims else None + ) self.launch_id += 1 return call(*args, launch_id=launch_id) @@ -343,19 +375,23 @@ class FfiKernel: assert num_inputs == self.num_inputs assert num_outputs == self.num_outputs - launch_bounds = launch_bounds_t(launch_desc.launch_dims) - # first kernel param is the launch bounds kernel_params = (ctypes.c_void_p * (1 + self.num_kernel_args))() - kernel_params[0] = ctypes.addressof(launch_bounds) - arg_refs = [] + batch_size = None # input and in-out args for i, input_arg in enumerate(self.input_args): if input_arg.is_array: buffer = inputs[i].contents - shape = buffer.dims[: input_arg.type.ndim] + shape = buffer.dims[: buffer.rank - input_arg.dtype_ndim] + if buffer.rank > input_arg.jax_ndim: + # handle batching + shape = collapse_batch_dims(shape, input_arg.type.ndim) + if batch_size is None: + batch_size = compute_batch_size( + buffer.dims[: buffer.rank], buffer.rank - input_arg.jax_ndim + ) strides = strides_from_shape(shape, input_arg.type.dtype) arg = array_t(buffer.data, 0, input_arg.type.ndim, shape, strides) kernel_params[i + 1] = ctypes.addressof(arg) @@ -370,12 +406,34 @@ class FfiKernel: # pure output args (skip in-out FFI buffers) for i, output_arg in enumerate(self.output_args): buffer = outputs[i + self.num_in_out].contents - shape = buffer.dims[: output_arg.type.ndim] + shape = buffer.dims[: buffer.rank - output_arg.dtype_ndim] + if buffer.rank > output_arg.jax_ndim: + # handle batching + shape = collapse_batch_dims(shape, output_arg.type.ndim) + if batch_size is None: + batch_size = compute_batch_size( + buffer.dims[: buffer.rank], buffer.rank - output_arg.jax_ndim + ) strides = strides_from_shape(shape, output_arg.type.dtype) arg = array_t(buffer.data, 0, output_arg.type.ndim, shape, strides) kernel_params[num_inputs + i + 1] = ctypes.addressof(arg) arg_refs.append(arg) # keep a reference + # determine launch bounds + if launch_desc.launch_dims is None: + # infer launch dims from argument shape, works with vmap + arr = arg_refs[self.first_array_arg] + launch_dims = arr.shape[: arr.ndim] + else: + # use specified launch dims + launch_dims = launch_desc.launch_dims + if batch_size is not None: + # roll batch size into the first launch dimension + launch_dims = (batch_size * launch_dims[0], *launch_dims[1:]) + + launch_bounds = launch_bounds_t(launch_dims) + 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) @@ -808,7 +866,7 @@ class FfiCallable: for i, arg in enumerate(self.input_args): if arg.is_array: buffer = inputs[i].contents - shape = buffer.dims[: buffer.rank - arg.dtype_ndim] + 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: @@ -819,7 +877,7 @@ class FfiCallable: # 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 = buffer.dims[: buffer.rank - arg.dtype_ndim] + 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) @@ -1095,6 +1153,7 @@ def jax_kernel( in_out_argnames=None, module_preload_mode=ModulePreloadMode.CURRENT_DEVICE, enable_backward: bool = False, + has_side_effect: bool = False, ): """Create a JAX callback from a Warp kernel. @@ -1103,21 +1162,23 @@ def jax_kernel( Args: kernel: The Warp kernel to launch. num_outputs: Specify the number of output arguments if greater than 1. - This must include the number of ``in_out_arguments``. + This must include the number of ``in_out_arguments``. vmap_method: String specifying how the callback transforms under ``vmap()``. - This argument can also be specified for individual calls. + This argument can also be specified for individual calls. launch_dims: Specify the default kernel launch dimensions. If None, launch - dimensions are inferred from the shape of the first array argument. - This argument can also be specified for individual calls. + dimensions are inferred from the shape of the first array argument. + This argument can also be specified for individual calls. output_dims: Specify the default dimensions of output arrays. If None, output - dimensions are inferred from the launch dimensions. - This argument can also be specified for individual calls. + dimensions are inferred from the launch dimensions. + This argument can also be specified for individual calls. in_out_argnames: Names of arguments that are both inputs and outputs (aliased buffers). These must be array arguments that appear before any pure output arguments in the kernel signature. The number of in-out arguments is included in ``num_outputs``. Not supported when ``enable_backward=True``. module_preload_mode: Specify the devices where the module should be preloaded. enable_backward: Enable automatic differentiation for this kernel. + has_side_effect: Whether the custom call has side effects. When True, + the FFI call will be executed even when the outputs are not used. Limitations: - All kernel arguments must be contiguous arrays or scalars. @@ -1129,21 +1190,41 @@ def jax_kernel( check_jax_version() + if isinstance(output_dims, dict): + hashable_output_dims = tuple(sorted(output_dims.items())) + elif hasattr(output_dims, "__len__"): + hashable_output_dims = tuple(output_dims) + else: + hashable_output_dims = output_dims + + if hasattr(launch_dims, "__len__"): + hashable_launch_dims = tuple(launch_dims) + else: + hashable_launch_dims = launch_dims + if not enable_backward: key = ( kernel.func, kernel.sig, num_outputs, vmap_method, - tuple(launch_dims) if launch_dims else launch_dims, - tuple(sorted(output_dims.items())) if output_dims else output_dims, + hashable_launch_dims, + hashable_output_dims, module_preload_mode, + has_side_effect, ) with _FFI_REGISTRY_LOCK: if key not in _FFI_KERNEL_REGISTRY: new_kernel = FfiKernel( - kernel, num_outputs, vmap_method, launch_dims, output_dims, in_out_argnames, module_preload_mode + kernel, + num_outputs, + vmap_method, + launch_dims, + output_dims, + in_out_argnames, + module_preload_mode, + has_side_effect=has_side_effect, ) _FFI_KERNEL_REGISTRY[key] = new_kernel @@ -1173,7 +1254,7 @@ def jax_kernel( static_args = [] for i, p in enumerate(parameters[:num_inputs]): param_type = p.annotation - if not isinstance(param_type, wp.array): + if not matches_array_class(param_type, wp.array): if param_type in wp._src.types.value_types: static_args.append(i) else: @@ -1183,7 +1264,7 @@ def jax_kernel( # determine launch dimensions from the shape of the first input array for i, p in enumerate(parameters[:num_inputs]): param_type = p.annotation - if isinstance(param_type, wp.array): + if matches_array_class(param_type, wp.array): arg = call_args[i] arg_shape = tuple(arg.shape) if hasattr(param_type.dtype, "_wp_scalar_type_"): @@ -1203,7 +1284,13 @@ def jax_kernel( fwd_kernel_wrapper.__annotations__ = {p.name: p.annotation for p in parameters} fwd_kernel_wrapper.__annotations__["return"] = None - jax_fwd_kernel = jax_callable(fwd_kernel_wrapper, num_outputs=num_outputs, vmap_method=vmap_method) + jax_fwd_kernel = jax_callable( + fwd_kernel_wrapper, + num_outputs=num_outputs, + vmap_method=vmap_method, + module_preload_mode=module_preload_mode, + has_side_effect=has_side_effect, + ) # backward arguments only include static args once bwd_arg_count = 2 * parameter_count - len(static_args) @@ -1285,6 +1372,8 @@ def jax_kernel( bwd_kernel_wrapper, num_outputs=len(bwd_input_params) - len(static_args), vmap_method=vmap_method, + module_preload_mode=module_preload_mode, + has_side_effect=has_side_effect, ) differentiable_input_indices = [i for i in range(num_inputs) if i not in static_args] @@ -1331,7 +1420,7 @@ def jax_kernel( if ann is None: continue # Check if annotation is a warp array type (annotation is an instance of wp.array) - is_array_ann = isinstance(ann, wp.array) + is_array_ann = matches_array_class(ann, wp.array) if not is_array_ann: continue dtype_ndim = 0 @@ -1355,6 +1444,15 @@ def jax_kernel( jax_func = jax.custom_vjp(jax_fwd_kernel, nondiff_argnums=tuple(static_args)) jax_func.defvjp(fwd_function, bwd_function) + key = ( + kernel.func, + kernel.sig, + num_outputs, + vmap_method, + module_preload_mode, + has_side_effect, + ) + if static_args: static_names = [parameters[i].name for i in static_args] @@ -1364,7 +1462,7 @@ def jax_kernel( _user_callable.__signature__ = signature # Cache differentiable wrapper - key = (kernel.func, kernel.sig, num_outputs, vmap_method, tuple(sorted(static_names))) + key = (*key, tuple(sorted(static_names))) with _FFI_REGISTRY_LOCK: cached = _FFI_DIFF_KERNEL_REGISTRY.get(key) if cached is None: @@ -1373,7 +1471,7 @@ def jax_kernel( return _FFI_DIFF_KERNEL_REGISTRY[key] # Cache differentiable wrapper (no static args) - key = (kernel.func, kernel.sig, num_outputs, vmap_method, ()) + key = (*key, ()) with _FFI_REGISTRY_LOCK: cached = _FFI_DIFF_KERNEL_REGISTRY.get(key) if cached is None: @@ -1426,6 +1524,8 @@ def jax_callable( graph_cache_max: Maximum number of cached graphs captured using ``GraphMode.WARP``. If ``None``, use ``warp.jax_experimental.get_jax_callable_default_graph_cache_max()``. module_preload_mode: Specify the devices where the module should be preloaded. + has_side_effect: Whether the custom call has side effects. When True, + the FFI call will be executed even when the outputs are not used. Limitations: - All kernel arguments must be contiguous arrays or scalars. @@ -1440,14 +1540,22 @@ def jax_callable( if graph_cache_max is None: graph_cache_max = FfiCallable.default_graph_cache_max + if isinstance(output_dims, dict): + hashable_output_dims = tuple(sorted(output_dims.items())) + elif hasattr(output_dims, "__len__"): + hashable_output_dims = tuple(output_dims) + else: + hashable_output_dims = output_dims + # Note: we don't include graph_cache_max in the key, it is applied below. key = ( func, num_outputs, graph_mode, vmap_method, - tuple(sorted(output_dims.items())) if output_dims else output_dims, + hashable_output_dims, module_preload_mode, + has_side_effect, ) with _FFI_REGISTRY_LOCK: diff --git a/mjx/pyproject.toml b/mjx/pyproject.toml index ff7133cd..60063b70 100644 --- a/mjx/pyproject.toml +++ b/mjx/pyproject.toml @@ -36,7 +36,7 @@ dependencies = [ [project.optional-dependencies] warp = [ - "warp-lang==1.11.1", + "warp-lang==1.12.0", ] [project.scripts]