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

Import NVIDIA/warp from GitHub.

Browse Source 
PiperOrigin-RevId: 892794388
Change-Id: Ic681071dc68bfaecaaab4e131c1c464dca99f9b5
This commit is contained in:
Taylor Howell
2026-04-01 01:53:33 -07:00
committed by Copybara-Service
parent 1640d6a4bc
commit 0b54638df4
4 changed files with 152 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/cuda_requirements.txt
+5 -5
View File
@@ -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
mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/custom_call.py
Vendored
+3 -3
View File
@@ -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.
mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/ffi.py
Vendored
+143 -35
View File
@@ -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:
mjx/pyproject.toml
+1 -1
View File
@@ -36,7 +36,7 @@ dependencies = [
[project.optional-dependencies]
warp = [
"warp-lang==1.11.1",
"warp-lang==1.12.0",
]
[project.scripts]