Import NVIDIA/warp from GitHub.

PiperOrigin-RevId: 892794388
Change-Id: Ic681071dc68bfaecaaab4e131c1c464dca99f9b5

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 \
+warp-lang==1.12.0 \
-    --hash=sha256:1ad11f1fa775269e991a3d55039152c8a504baf86701c849b485cb8e66c49d15 \
+    --hash=sha256:c78c3701d5cad86c30ef5017410d294ec46a396bb0d502ee1c98743494f3a62f \
-    --hash=sha256:8b098f41e71d421d80ee7562e38aa8380ff6b0d3b4c6ee866cfbdef733ac5bdc \
+    --hash=sha256:a1436f60a1881cd94f787e751a83fc0987626be2d3e2b4e74c64a6947c6d1266 \
-    --hash=sha256:5d0904b0eefcc81f39ba65375427a3de99006088aa43e24a9011263f07d0cd07 \
+    --hash=sha256:a2d6decba693aba5b828573c4414fd6a3f4c4a934db9c322736ef2b3fa99fe76 \
-    --hash=sha256:15dc10aa51fb0fdbe1ca16d52e5fadca35a47ffd9d0c636826506f96bb2e7c41
+    --hash=sha256:697248edd2f1e2952f50e3db33b214af76173641a8894aacc467bed6dc247f8a

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.

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_KERNEL_REGISTRY: dict[tuple, FfiKernel] = {}
-_FFI_DIFF_KERNEL_REGISTRY: dict[DiffKernelCacheKey, Callable] = {}
+_FFI_DIFF_KERNEL_REGISTRY: dict[tuple, Callable] = {}
-_FFI_CALLABLE_REGISTRY: dict[str, FfiCallable] = {}
+_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.
 
@@ -1118,6 +1177,8 @@ def jax_kernel(
             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,
+            hashable_launch_dims,
-            tuple(sorted(output_dims.items())) if output_dims else output_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

@@ -36,7 +36,7 @@ dependencies = [
 
 [project.optional-dependencies]
 warp = [
-    "warp-lang==1.11.1",
+    "warp-lang==1.12.0",
 ]
 
 [project.scripts]