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

Cache MJX-Warp FFI callables

Browse Source 
Equivalent JAX retraces rebuilt flattened wrapper functions, causing
Warp to register a new FFI callable for every trace.

Reuse callables with matching structural configurations while keeping
distinct shim functions isolated. This keeps Warp's registry and graph
cache effective without changing callback lifetime.

Document the structural key and lock scope so future changes preserve
the cache's intended boundaries.

Signed-off-by: Eric Shi <ershi@nvidia.com>
This commit is contained in:
Eric Shi
2026-08-05 20:13:35 +00:00
parent a1d772c9ad
commit ab4fe36121
2 changed files with 154 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/warp/ffi.py
+93 -27
View File
@@ -17,6 +17,7 @@
import dataclasses
import functools
import inspect
import threading
import typing
from typing import Any, Callable, Optional, Sequence, Tuple, Union
@@ -30,6 +31,15 @@ from mujoco.mjx._src.types import tree_path_to_attr_str
from mujoco.mjx.warp import types as mjx_warp_types
# ``warp_ffi.jax_callable`` keys its registry by wrapper function and
# configuration. Cache generated wrappers here by the original shim and MJX's
# flattened call structure so equivalent retraces reuse the same Warp entry.
_JAX_CALLABLE_VARIADIC_TUPLE_REGISTRY: dict[
tuple[Any, ...], Callable[..., Any]
] = {}
_JAX_CALLABLE_VARIADIC_TUPLE_REGISTRY_LOCK = threading.Lock()
def flatten_signature(signature: inspect.Signature, args: Tuple[Any, ...]):
"""Flattens a tuple/dataclass signature."""
@@ -109,38 +119,94 @@ def jax_callable_variadic_tuple(
):
"""Wraps a JAX callable to support variadic tuples and dataclasses."""
# Snapshot mapping and name-set options into stable cache-key forms. Warp
# consumes them by argument name, so caller ordering does not distinguish
# callable configurations.
hashable_output_dims = (
None if output_dims is None else tuple(sorted(output_dims.items()))
)
hashable_in_out_argnames = (
tuple(sorted(in_out_argnames)) if in_out_argnames else None
)
hashable_stage_in_argnames = (
tuple(sorted(stage_in_argnames)) if stage_in_argnames else None
)
hashable_stage_out_argnames = (
tuple(sorted(stage_out_argnames)) if stage_out_argnames else None
)
def callable_wrapper(*args, **kwargs):
def func_wrapper(*flat_args, **kwargs):
num_inputs = in_tree.num_leaves
flat_inputs = flat_args[:num_inputs]
output_buffers = flat_args[num_inputs:]
unflat_args = jax.tree.unflatten(in_tree, flat_inputs)
return func(*unflat_args, *output_buffers, **kwargs)
# Provide a flattened signature for the Warp callable machinery.
flat_args, in_tree = jax.tree.flatten(args)
new_signature = flatten_signature(inspect.signature(func), args)
func_wrapper.__signature__ = new_signature # pyrefly: ignore[missing-attribute]
func_wrapper.__annotations__ = {
p.name: p.annotation
for p in new_signature.parameters.values()
if p.annotation is not inspect.Parameter.empty
}
if new_signature.return_annotation is not inspect.Signature.empty:
func_wrapper.__annotations__['return'] = new_signature.return_annotation
my_callable = warp_ffi.jax_callable(
func_wrapper,
num_outputs=num_outputs,
graph_mode=graph_mode,
vmap_method=vmap_method,
output_dims=output_dims,
in_out_argnames=in_out_argnames,
stage_in_argnames=stage_in_argnames,
stage_out_argnames=stage_out_argnames,
has_side_effect=has_side_effect,
# Cache the wrapper's structural ABI, not per-call leaves. The flattened
# signature defines Warp's arguments and the PyTree defines reconstruction.
# Leaf arrays and tracers are forwarded on every invocation; keying on them
# would prevent equivalent retraces from reusing the same FFI target.
callable_cache_key = (
func,
new_signature,
in_tree,
num_outputs,
graph_mode,
vmap_method,
hashable_output_dims,
hashable_in_out_argnames,
hashable_stage_in_argnames,
hashable_stage_out_argnames,
has_side_effect,
)
flat_args, in_tree = jax.tree.flatten(args)
# Serialize construction so concurrent traces cannot register duplicate
# Warp callbacks for the same structural key.
with _JAX_CALLABLE_VARIADIC_TUPLE_REGISTRY_LOCK:
my_callable = _JAX_CALLABLE_VARIADIC_TUPLE_REGISTRY.get(
callable_cache_key
)
if my_callable is None:
# Restore the original PyTree inputs; Warp appends output buffers.
def func_wrapper(*flat_args, **kwargs):
num_inputs = in_tree.num_leaves
flat_inputs = flat_args[:num_inputs]
output_buffers = flat_args[num_inputs:]
unflat_args = jax.tree.unflatten(in_tree, flat_inputs)
return func(*unflat_args, *output_buffers, **kwargs)
# Warp derives the FFI ABI from this synthetic signature and
# annotations.
func_wrapper.__signature__ = ( # pyrefly: ignore[missing-attribute]
new_signature
)
func_wrapper.__annotations__ = {
p.name: p.annotation
for p in new_signature.parameters.values()
if p.annotation is not inspect.Parameter.empty
}
if new_signature.return_annotation is not inspect.Signature.empty:
func_wrapper.__annotations__['return'] = (
new_signature.return_annotation
)
# Constructing the callable registers its FFI target.
my_callable = warp_ffi.jax_callable(
func_wrapper,
num_outputs=num_outputs,
graph_mode=graph_mode,
vmap_method=vmap_method,
output_dims=(
None
if hashable_output_dims is None
else dict(hashable_output_dims)
),
in_out_argnames=hashable_in_out_argnames,
stage_in_argnames=hashable_stage_in_argnames,
stage_out_argnames=hashable_stage_out_argnames,
has_side_effect=has_side_effect,
)
_JAX_CALLABLE_VARIADIC_TUPLE_REGISTRY[callable_cache_key] = my_callable
# Invocation may re-enter JAX or Warp and needs no registry serialization.
return my_callable(*flat_args, **kwargs)
return callable_wrapper
mjx/mujoco/mjx/warp/ffi_test.py
+61
View File
@@ -0,0 +1,61 @@
# Copyright 2026 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for MJX-Warp FFI helpers."""
from unittest import mock
from absl.testing import absltest
from mujoco.mjx.warp import ffi
class FfiTest(absltest.TestCase):
def test_jax_callable_variadic_tuple_cache(self):
def func_a(values: tuple[int, ...], output: int):
del values, output
def func_b(values: tuple[int, ...], output: int):
del values, output
# Count callback registrations without creating real JAX FFI targets.
def create_callable(*args, **kwargs):
del args, kwargs
return mock.Mock()
with mock.patch.object(ffi, '_JAX_CALLABLE_VARIADIC_TUPLE_REGISTRY', {}):
with mock.patch.object(
ffi.warp_ffi, 'jax_callable', side_effect=create_callable
) as jax_callable:
wrapper_a = ffi.jax_callable_variadic_tuple(func_a)
wrapper_a_again = ffi.jax_callable_variadic_tuple(func_a)
wrapper_b = ffi.jax_callable_variadic_tuple(func_b)
# Equivalent traces of the same callable reuse one callback.
wrapper_a((1, 2))
wrapper_a_again((3, 4))
self.assertEqual(jax_callable.call_count, 1)
# Distinct shim functions with the same signature do not share a target.
wrapper_b((5, 6))
self.assertEqual(jax_callable.call_count, 2)
# A different tuple arity changes the signature and PyTree cache key.
wrapper_b((1, 2, 3))
self.assertEqual(jax_callable.call_count, 3)
if __name__ == '__main__':
absltest.main()