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

Import NVIDIA/warp 1.13.0 from GitHub.

Browse Source 
PiperOrigin-RevId: 923878216
Change-Id: If9df5e0672e956eb62874fb6f4f74ba3a02977b3
This commit is contained in:
Taylor Howell
2026-05-30 04:47:29 -07:00
committed by Copybara-Service
parent 3a6e0ee9b3
commit b9c7a4b81a
5 changed files with 29 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/cuda_requirements.txt
+5 -5
View File
@@ -20,8 +20,8 @@ jax-cuda12-pjrt==0.8.3; python_version >= '3.13' \
jax-cuda12-pjrt==0.5.3; python_version < '3.13' \
--hash=sha256:04ee111eaf5fc2692978ad4a5c84d5925e42eb05c1701849ba3a53f6515400cc \
--hash=sha256:c5378306568ba0c81b230a779dd3194c9dd10339ab6360ae80928108d37e7f75
warp-lang==1.12.1 \
--hash=sha256:98df3533a6c40a33cce961f8efa991006b30c9d286356e4cd77ea8ce86928f1d \
--hash=sha256:6bf01f10509488ba8eacaf4ec7fcf7cfbd503118b22e002ecba407b40a17424e \
--hash=sha256:af6d680e79c1be6e46ddf80ecaa358f222804f882f4683260a7b4abd80a0981b \
--hash=sha256:826b2f93df8e47eac0c751a8eb5a0533e2fc5434158c8896a63be53bfbd728c7
warp-lang==1.13.0 \
--hash=sha256:4375f572301991fe0fbf0af29fc84d76cd27d531432d6df8452b25088c21ea5a \
--hash=sha256:ac2479c70ad410d58deb088c2a64168792be588efc1bec22dc51f92238e8c8c3 \
--hash=sha256:476b54f0dcf6767f23305a328660a9a74d01e371b6b7c3d77e03e18b4a1bb1a5 \
--hash=sha256:47975ea07252d45a4d09d2d1a6cffc55002fa7fbde771c8dceb1baf4b32d4fb8
mjx/mujoco/mjx/_src/io.py
+2 -2
View File
@@ -179,8 +179,8 @@ def _wp_to_np_type(wp_field: Any, name: str = '') -> Any:
wp_field.shape = wp_field.shape[1:]
# warp scalars
wp_dtype = type(wp_field)
if wp_dtype in wp.types.warp_type_to_np_dtype:
return wp.types.warp_type_to_np_dtype[wp_dtype](wp_field)
if wp_dtype in wp._src.types.warp_type_to_np_dtype:
return wp.dtype_to_numpy(wp_dtype)(wp_field)
# warp arrays
if isinstance(wp_field, wp.array):
mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/custom_call.py
Vendored
+7 -5
View File
@@ -131,13 +131,13 @@ def _warp_custom_callback(stream, buffers, opaque, opaque_len):
# Launch the kernel.
wp._src.context.runtime.core.wp_cuda_launch_kernel(
device.context, hooks.forward, bounds.size, 0, 256, hooks.forward_smem_bytes, kernel_params, stream
device.context, hooks.forward, bounds.size, 0, 256, hooks.forward_smem_bytes, kernel_params, stream, None
)
def _create_jax_warp_primitive():
import jax # noqa: PLC0415
from jax._src.interpreters import batching # noqa: PLC0415
from jax.interpreters import batching # noqa: PLC0415
from jax.interpreters import mlir # noqa: PLC0415
from jax.interpreters.mlir import ir # noqa: PLC0415
from jaxlib.hlo_helpers import custom_call # noqa: PLC0415
@@ -170,9 +170,9 @@ def _create_jax_warp_primitive():
# Figure out the number of outputs.
wp_kernel = _registered_kernels[params["kernel"]]
output_count = len(wp_kernel.adj.args) - len(args)
shape, dim = next((a.shape, d) for a, d in zip(args, dims) if d is not None)
shape, dim = next((a.shape, d) for a, d in zip(args, dims, strict=True) if d is not None)
size = shape[dim]
args = [batching.bdim_at_front(a, d, size) if len(a.shape) else a for a, d in zip(args, dims)]
args = [batching.bdim_at_front(a, d, size) if len(a.shape) else a for a, d in zip(args, dims, strict=True)]
# Create the batched primitive.
return _jax_warp_p.bind(*args, **params), [dims[0]] * output_count
@@ -211,6 +211,7 @@ def _create_jax_warp_primitive():
def base_type_to_jax_ir(warp_dtype):
warp_to_jax_dict = {
wp.float16: ir.F16Type.get(),
wp.bfloat16: ir.BF16Type.get(),
wp.float32: ir.F32Type.get(),
wp.float64: ir.F64Type.get(),
wp.int8: ir.IntegerType.get_signless(8),
@@ -232,6 +233,7 @@ def _create_jax_warp_primitive():
def base_type_is_compatible(warp_type, jax_ir_type):
jax_ir_to_warp = {
"f16": wp.float16,
"bf16": wp.bfloat16,
"f32": wp.float32,
"f64": wp.float64,
"i8": wp.int8,
@@ -324,7 +326,7 @@ def _create_jax_warp_primitive():
# Figure out the types and shapes of the input arrays.
arg_strings = []
operand_layouts = []
for actual, warg in zip(args, wp_kernel.adj.args):
for actual, warg in zip(args, wp_kernel.adj.args, strict=False):
wtype = warg.type
rtt = ir.RankedTensorType(actual.type)
mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/ffi.py
Vendored
+14 -13
View File
@@ -8,8 +8,8 @@ import ctypes
import inspect
import threading
import traceback
from collections.abc import Callable
from enum import IntEnum
from typing import Callable
import jax
@@ -25,6 +25,7 @@ from warp._src.types import (
type_size_in_bytes,
type_to_warp,
)
from warp._src.utils import warn
from .xla_ffi import *
@@ -249,19 +250,18 @@ 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}',"
f" expected {input_arg.jax_scalar_type}, got {input_value.dtype}"
f"Invalid data type for array argument '{input_arg.name}', 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
@@ -455,6 +455,7 @@ class FfiKernel:
hooks.forward_smem_bytes,
kernel_params,
stream,
None, # apic_info
)
else:
wp._src.context.runtime.core.wp_cpu_launch_kernel(
@@ -467,9 +468,7 @@ class FfiKernel:
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}"
)
@@ -617,7 +616,9 @@ class FfiCallable:
self.input_output_aliases = input_output_aliases
# register the callback
FFI_CCALLFUNC = ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.POINTER(XLA_FFI_CallFrame))
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)
@@ -1071,7 +1072,7 @@ class FfiCallable:
call_desc.capture = capture
else:
# not capturing or on CPU
# not capturing
self.func(*arg_list)
except Exception as e:
@@ -1359,7 +1360,7 @@ def jax_kernel(
try:
gi.zero_()
except Exception as e:
wp.utils.warn(f"Failed to zero gradient array: {e}", stacklevel=2)
warn(f"Failed to zero gradient array: {e}", stacklevel=2)
raise e
# NOTE: We cannot use a passed launch_dims here, the backward rule doesn't receive it (and it could be wrong under pmap/vmap).
@@ -1445,7 +1446,7 @@ def jax_kernel(
non_static_inputs, output_vals_tuple = residuals
input_vals = list(non_static_inputs)
for i, v in zip(static_args, nondiff_vals):
for i, v in zip(static_args, nondiff_vals, strict=True):
input_vals.insert(i, v)
# Normalize grad outputs and handle nested containers (e.g., single tuple for multi-output)
@@ -1461,7 +1462,7 @@ def jax_kernel(
out_dims_map = {}
param_ann = {p.name: p.annotation for p in parameters[:num_inputs]}
for name, val in zip(differentiable_input_names, non_static_inputs):
for name, val in zip(differentiable_input_names, non_static_inputs, strict=True):
ann = param_ann.get(name)
if ann is None:
continue
mjx/pyproject.toml
+1 -1
View File
@@ -37,7 +37,7 @@ dependencies = [
[project.optional-dependencies]
warp = [
"warp-lang==1.12.1",
"warp-lang==1.13.0",
]
dev = [
"isort",