Import NVIDIA/warp from GitHub.

PiperOrigin-RevId: 872969263
Change-Id: Ic890cb53c32593c158a1f3926906b492c7f85348

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.0 \
-    --hash=sha256:3a4f1c9a6e721d7de7d6dad6b242c54afaf20c6e14a767c0da03e5e963fcc13c \
-    --hash=sha256:524dce20de6162ba25333552168ebf430973050e00d9f8116b8df41a60d25d6e \
-    --hash=sha256:1ae6cfc226107f96e4d495b41a3dab32488e8ee8f074b0e1bcaf22e7fb8c904d \
-    --hash=sha256:80d8493cbe243a3510134f3af289646d7bd7484217a30ecf565d676466ef8a5e
+warp-lang==1.11.1 \
+    --hash=sha256:1ad11f1fa775269e991a3d55039152c8a504baf86701c849b485cb8e66c49d15 \
+    --hash=sha256:8b098f41e71d421d80ee7562e38aa8380ff6b0d3b4c6ee866cfbdef733ac5bdc \
+    --hash=sha256:5d0904b0eefcc81f39ba65375427a3de99006088aa43e24a9011263f07d0cd07 \
+    --hash=sha256:15dc10aa51fb0fdbe1ca16d52e5fadca35a47ffd9d0c636826506f96bb2e7c41

mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/ffi.py

@@ -62,11 +62,22 @@ def check_jax_version():
 
 
 class GraphMode(IntEnum):
-    NONE = 0  # don't capture a graph
-    JAX = 1  # let JAX capture a graph
-    WARP = 2  # let Warp capture a graph
-    WARP_STAGED = 3  # use Warp graph with staging buffers, copy inside of the graph
-    WARP_STAGED_EX = 4  # use Warp graph with staging buffers, copy outside of the graph
+    """CUDA graph capture modes for :func:`warp.jax_experimental.jax_callable`.
+
+    These modes control whether JAX or Warp captures a CUDA graph, and whether
+    staging buffers are used when capturing with Warp.
+    """
+
+    NONE = 0
+    """Disable graph capture. Use when operations are not CUDA-graph compatible (for example, host synchronization)."""
+    JAX = 1
+    """Let JAX capture the graph so the callable can be used as a subgraph within a larger JAX capture."""
+    WARP = 2
+    """Let Warp capture the graph and replay it for matching buffer addresses."""
+    WARP_STAGED = 3
+    """Capture a Warp graph using staging buffers and insert memcpy nodes inside the graph."""
+    WARP_STAGED_EX = 4
+    """Capture a Warp graph using staging buffers and perform memcpy outside the graph."""
 
 
 class ModulePreloadMode(IntEnum):
@@ -682,12 +693,13 @@ class FfiCallable:
                 assert num_outputs == self.num_outputs
 
                 cuda_stream = get_stream_from_callframe(call_frame.contents)
+                device_ordinal = get_device_ordinal_from_callframe(call_frame.contents)
 
                 if self.graph_mode == GraphMode.WARP:
                     # check if we already captured an identical call
                     ip = [inputs[i].contents.data for i in self.array_input_indices]
                     op = [outputs[i].contents.data for i in self.array_output_indices]
-                    capture_key = hash((call_id, *ip, *op))
+                    capture_key = hash((device_ordinal, call_id, *ip, *op))
                     capture = self.captures.get(capture_key)
 
                     # launch existing graph

mjx/pyproject.toml

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