Skip to content

Add JAX FFI Host support #1446

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
loney7 wants to merge 1 commit into
base: main
Choose a base branch
from
+480 −62
Open

Add JAX FFI Host support #1446

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,8 @@
host-return form of `wp.utils.runlength_encode()`, and `wp.copy()` other than a contiguous same-device copy
(cross-device or non-contiguous copies) cannot be captured on CUDA and raise during APIC capture
([GH-1431](https://github.com/NVIDIA/warp/issues/1431)).
- Add support for running JAX FFI callbacks on the CPU (Host) platform in addition to CUDA
([GH-1446](https://github.com/NVIDIA/warp/issues/1446)).

### Removed

Expand Down
191 changes: 132 additions & 59 deletions warp/_src/jax/ffi.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -235,10 +235,22 @@ def __init__(
# register the callback
jax = _get_jax()
FFI_CCALLFUNC = ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.POINTER(XLA_FFI_CallFrame))
self.callback_func = FFI_CCALLFUNC(self.ffi_callback)
ffi_ccall_address = ctypes.cast(self.callback_func, ctypes.c_void_p)
ffi_capsule = jax.ffi.pycapsule(ffi_ccall_address.value)
jax.ffi.register_ffi_target(self.name, ffi_capsule, platform="CUDA")

def make_callback(platform):
def callback(call_frame):
return self.ffi_callback(call_frame, platform)

return callback

self.callback_func_cuda = FFI_CCALLFUNC(make_callback("CUDA"))
ffi_capsule_cuda = jax.ffi.pycapsule(ctypes.cast(self.callback_func_cuda, ctypes.c_void_p).value)
jax.ffi.register_ffi_target(self.name, ffi_capsule_cuda, platform="CUDA")
try:
self.callback_func_host = FFI_CCALLFUNC(make_callback("Host"))
ffi_capsule_host = jax.ffi.pycapsule(ctypes.cast(self.callback_func_host, ctypes.c_void_p).value)
jax.ffi.register_ffi_target(self.name, ffi_capsule_host, platform="Host")
except AttributeError:
pass

def __call__(self, *args, output_dims=None, launch_dims=None, vmap_method=None):
jax = _get_jax()
Expand Down Expand Up @@ -332,17 +344,19 @@ def __call__(self, *args, output_dims=None, launch_dims=None, vmap_method=None):
# preload on the specified devices
if self.module_preload_mode == JaxModulePreloadMode.CURRENT_DEVICE:
device = wp.device_from_jax(get_jax_device())
self.kernel.module.load(device)
block_dim = 256 if device.is_cuda else 1
self.kernel.module.load(device, block_dim)
elif self.module_preload_mode == JaxModulePreloadMode.ALL_DEVICES:
for d in jax.local_devices():
try:
dev = wp.device_from_jax(d)
except Exception:
# ignore unsupported devices like TPUs
pass
# we only support CUDA devices for now
if dev.is_cuda:
self.kernel.module.load(dev)
# we only support CUDA and CPU devices for now
if dev.is_cuda or dev.is_cpu:
block_dim = 256 if dev.is_cuda else 1
self.kernel.module.load(dev, block_dim)

# save launch data to be retrieved by callback
launch_id = self.launch_id
Expand All @@ -353,7 +367,7 @@ def __call__(self, *args, output_dims=None, launch_dims=None, vmap_method=None):

return call(*args, launch_id=launch_id)

def ffi_callback(self, call_frame):
def ffi_callback(self, call_frame, platform):
try:
# On the first call, XLA runtime will query the API version and traits
# metadata using the |extension| field. Let us respond to that query
Expand All @@ -365,10 +379,11 @@ def ffi_callback(self, call_frame):
metadata_ext = ctypes.cast(extension, ctypes.POINTER(XLA_FFI_Metadata_Extension))
metadata_ext.contents.metadata.contents.api_version.major_version = 0
metadata_ext.contents.metadata.contents.api_version.minor_version = 1
# Turn on CUDA graphs for this handler.
metadata_ext.contents.metadata.contents.traits = (
XLA_FFI_Handler_TraitsBits.COMMAND_BUFFER_COMPATIBLE
)
if platform == "CUDA":
# Turn on CUDA graphs for this handler.
metadata_ext.contents.metadata.contents.traits = (
XLA_FFI_Handler_TraitsBits.COMMAND_BUFFER_COMPATIBLE
)
return None

# Lock is required to prevent race conditions when callback is invoked
Expand Down Expand Up @@ -448,32 +463,53 @@ def ffi_callback(self, call_frame):
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)

# get kernel hooks
hooks = self.kernel.module.get_kernel_hooks(self.kernel, device)
if platform == "CUDA":
device = wp.get_cuda_device(get_device_ordinal_from_callframe(call_frame.contents))
stream = get_stream_from_callframe(call_frame.contents)
else:
device = wp.get_device("cpu")
stream = None

# get kernel hooks (ensure module is loaded on the target device,
# e.g. Host callbacks may run on CPU even when preloaded on CUDA).
# CPU tile execution is single-lane, so use block_dim=1 to compile
# WP_TILE_BLOCK_DIM correctly for explicitly shaped tile operations.
block_dim = 256 if device.is_cuda else 1
module_exec = self.kernel.module.load(device, block_dim)
hooks = module_exec.get_kernel_hooks(self.kernel)
assert hooks.forward, "Failed to find kernel entry point"

# reject non-cluster-aligned grids with a clear Python error instead
# of a cryptic native CUDA error (block_dim=256, max_blocks=0 below)
_validate_cluster_launch(hooks.cluster_dim, launch_bounds.size, 256, 0)

# launch the kernel (cluster_dim is cached on the hooks at load time)
if wp._src.context.runtime.core.wp_cuda_launch_kernel(
device.context,
hooks.forward,
launch_bounds.size,
0,
256,
int(self.kernel.grid_stride),
hooks.cluster_dim,
hooks.forward_smem_bytes,
kernel_params,
stream,
None, # apic_info
):
_raise_cuda_launch_error(self.kernel, device)
if platform == "CUDA":
if wp._src.context.runtime.core.wp_cuda_launch_kernel(
device.context,
hooks.forward,
launch_bounds.size,
0,
256,
int(self.kernel.grid_stride),
hooks.cluster_dim,
hooks.forward_smem_bytes,
kernel_params,
stream,
None, # apic_info
):
_raise_cuda_launch_error(self.kernel, device)
else:
args_struct, _ = wp._src.context._build_cpu_args_structs(
self.kernel, hooks, [launch_bounds, *arg_refs], False
)
wp._src.context.runtime.core.wp_cpu_launch_kernel(
ctypes.cast(hooks.forward, ctypes.c_void_p),
ctypes.byref(launch_bounds),
ctypes.byref(args_struct) if args_struct is not None else None,
None,
None,
)

except Exception as e:
print(traceback.format_exc())
Expand Down Expand Up @@ -628,10 +664,22 @@ def __init__(
# register the callback
jax = _get_jax()
FFI_CCALLFUNC = ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.POINTER(XLA_FFI_CallFrame))
self.callback_func = FFI_CCALLFUNC(self.ffi_callback)
ffi_ccall_address = ctypes.cast(self.callback_func, ctypes.c_void_p)
ffi_capsule = jax.ffi.pycapsule(ffi_ccall_address.value)
jax.ffi.register_ffi_target(self.name, ffi_capsule, platform="CUDA")

def make_callback(platform):
def callback(call_frame):
return self.ffi_callback(call_frame, platform)

return callback

self.callback_func_cuda = FFI_CCALLFUNC(make_callback("CUDA"))
ffi_capsule_cuda = jax.ffi.pycapsule(ctypes.cast(self.callback_func_cuda, ctypes.c_void_p).value)
jax.ffi.register_ffi_target(self.name, ffi_capsule_cuda, platform="CUDA")
try:
self.callback_func_host = FFI_CCALLFUNC(make_callback("Host"))
ffi_capsule_host = jax.ffi.pycapsule(ctypes.cast(self.callback_func_host, ctypes.c_void_p).value)
jax.ffi.register_ffi_target(self.name, ffi_capsule_host, platform="Host")
except AttributeError:
pass

def __call__(self, *args, output_dims=None, vmap_method=None):
jax = _get_jax()
Expand Down Expand Up @@ -715,25 +763,27 @@ def __call__(self, *args, output_dims=None, vmap_method=None):
module = wp.get_module(self.func.__module__)
if self.module_preload_mode == JaxModulePreloadMode.CURRENT_DEVICE:
device = wp.device_from_jax(get_jax_device())
module.load(device)
block_dim = 256 if device.is_cuda else 1
module.load(device, block_dim)
elif self.module_preload_mode == JaxModulePreloadMode.ALL_DEVICES:
for d in jax.local_devices():
try:
dev = wp.device_from_jax(d)
except Exception:
# ignore unsupported devices like TPUs
pass
# we only support CUDA devices for now
if dev.is_cuda:
module.load(dev)
# we only support CUDA and CPU devices for now
if dev.is_cuda or dev.is_cpu:
block_dim = 256 if dev.is_cuda else 1
module.load(dev, block_dim)

# save call data to be retrieved by callback
call_id = self.call_id
self.call_descriptors[call_id] = FfiCallDesc(static_inputs)
self.call_id += 1
return call(*args, call_id=call_id)

def ffi_callback(self, call_frame):
def ffi_callback(self, call_frame, platform):
try:
# On the first call, XLA runtime will query the API version and traits
# metadata using the |extension| field. Let us respond to that query
Expand Down Expand Up @@ -773,8 +823,12 @@ def ffi_callback(self, call_frame):
assert num_inputs == self.num_inputs
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 platform == "CUDA":
cuda_stream = get_stream_from_callframe(call_frame.contents)
device_ordinal = get_device_ordinal_from_callframe(call_frame.contents)
else:
cuda_stream = None
device_ordinal = None

if self.graph_mode == JaxCallableGraphMode.WARP:
# check if we already captured an identical call
Expand Down Expand Up @@ -878,9 +932,12 @@ def ffi_callback(self, call_frame):
# early out
return

device_ordinal = get_device_ordinal_from_callframe(call_frame.contents)
device = wp.get_cuda_device(device_ordinal)
stream = wp.Stream(device, cuda_stream=cuda_stream)
if platform == "CUDA":
device = wp.get_cuda_device(device_ordinal)
stream = wp.Stream(device, cuda_stream=cuda_stream)
else:
device = wp.get_device("cpu")
Comment thread
greptile-apps[bot] marked this conversation as resolved.
stream = None

# reconstruct the argument list
arg_list = []
Expand All @@ -906,7 +963,7 @@ def ffi_callback(self, call_frame):

# call the Python function with reconstructed arguments
with wp.ScopedStream(stream, sync_enter=False):

@shi-eric shi-eric Jul 6, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

On the Host platform stream is None here, and ScopedStream(None) is a no-op (its
__enter__/__exit__ do nothing when the stream is None), so no device scope is
established around the callback. Device-less wp.launch() calls inside the callback then
resolve to Warp's default device, which is cuda:0 on any machine with a visible GPU. The
kernels launch asynchronously on CUDA with CPU buffer pointers, producing silently wrong
output (or an async CUDA error 700: illegal memory access, depending on whether the GPU
can access pageable host memory).

This is the failure on the GPU runners in
test-warp-gpu-linux
and
test-warp-gpu-linux-mgpu-python310:
test_ffi_jax_callable_host_in_out fails with stale/partially written output (8 of 1M
elements stale in one job, 100% in the other). The other host jax_callable tests only
pass by timing luck; the same pattern as scale_func produced wrong results on my machine.

Suggested change
with wp.ScopedStream(stream, sync_enter=False):
# ScopedStream makes the stream's device current, but there is
# no stream on the Host platform, so use ScopedDevice to route
# device-less launches in the callback to the CPU
if stream is not None:
exec_scope = wp.ScopedStream(stream, sync_enter=False)
else:
exec_scope = wp.ScopedDevice(device)
with exec_scope:

Verified locally on 2x RTX 3090 (CPU jaxlib and jax[cuda13]): before this change, the
host jax_callable tests fail with GPUs visible and pass with CUDA_VISIBLE_DEVICES=;
after it, all host and CUDA FFI tests pass. The CUDA branch is unchanged.

A regression test can pin this down deterministically by asserting on the resolved device
instead of comparing outputs (output comparisons can pass by timing luck on drivers where
the GPU can dereference pageable host memory):

@unittest.skipUnless(_jax_version() >= (0, 5, 0), "Jax version too old")
@unittest.skipUnless(wp.is_cpu_available(), "Requires CPU backend")
def test_ffi_jax_callable_host_default_device(test, device):
    """Verify that Host callbacks route device-less launches to the CPU.

    Regression test for the Host (CPU) FFI path: the callback must run with
    the CPU device current so that ``wp.launch()`` calls without an explicit
    device run on the CPU even when Warp's default device is a CUDA device.
    """
    import jax.numpy as jp  # noqa: PLC0415

    jax_callable = wp.jax_callable

    # devices resolved by device-less launches inside the callback
    resolved_devices = []

    def record_device_func(a: wp.array[float], b: wp.array[float]):
        resolved_devices.append(wp.get_device())
        wp.launch(scale_kernel, dim=a.shape, inputs=[a, 2.0], outputs=[b])

    jax_func = jax_callable(record_device_func)

    @jax.jit
    def f():
        a = jp.arange(ARRAY_SIZE, dtype=jp.float32)
        return jax_func(a, output_dims={"b": a.shape})

    # force a non-CPU default device to expose launches that don't inherit the
    # callback's device scope
    device_scope = wp.ScopedDevice("cuda:0") if wp.is_cuda_available() else contextlib.nullcontext()

    with device_scope, jax.default_device(jax.devices("cpu")[0]):
        (b,) = f()
        jax.block_until_ready(b)

    test.assertGreater(len(resolved_devices), 0)
    for d in resolved_devices:
        test.assertEqual(d, wp.get_device("cpu"))

    assert_np_equal(np.asarray(b), 2 * np.arange(ARRAY_SIZE, dtype=np.float32))

(plus an add_function_test(TestJax, "test_ffi_jax_callable_host_default_device", ..., devices=None)
registration next to the other host tests)

if stream.is_capturing:
if stream is not None and stream.is_capturing:
# capturing with JAX
with wp.ScopedCapture(external=True) as capture:
self.func(*arg_list)
Expand Down Expand Up @@ -1212,7 +1269,8 @@ def jax_kernel(
- Scalars must be static arguments in JAX.
- Input and input-output arguments must precede the output arguments in the ``kernel`` definition.
- There must be at least one output or input-output argument.
- Only the CUDA backend is supported.
- Only the CUDA and CPU (Host) backends are supported.
- ``enable_backward`` is not supported on the CPU (Host) backend.
- ``output_dims`` and ``in_out_argnames`` are not supported when ``enable_backward=True``.
"""

Expand Down Expand Up @@ -1593,7 +1651,8 @@ def jax_callable(
- Scalars must be static arguments in JAX.
- Input and input-output arguments must precede the output arguments in the ``func`` definition.
- There must be at least one output or input-output argument.
- Only the CUDA backend is supported.
- Only the CUDA and CPU (Host) backends are supported.
- CUDA graph modes are not supported on the CPU (Host) backend.
"""

check_jax_version()
Expand Down Expand Up @@ -1692,7 +1751,7 @@ def register_ffi_callback(name: str, func: Callable, graph_compatible: bool = Tr

# TODO check that the name is not already registered

def ffi_callback(call_frame):
def ffi_callback(call_frame, platform):
try:
extension = call_frame.contents.extension_start
# On the first call, XLA runtime will query the API version and traits
Expand All @@ -1704,7 +1763,7 @@ def ffi_callback(call_frame):
metadata_ext = ctypes.cast(extension, ctypes.POINTER(XLA_FFI_Metadata_Extension))
metadata_ext.contents.metadata.contents.api_version.major_version = 0
metadata_ext.contents.metadata.contents.api_version.minor_version = 1
if graph_compatible:
if platform == "CUDA" and graph_compatible:
# Turn on CUDA graphs for this handler.
metadata_ext.contents.metadata.contents.traits = (
XLA_FFI_Handler_TraitsBits.COMMAND_BUFFER_COMPATIBLE
Expand All @@ -1718,13 +1777,13 @@ def ffi_callback(call_frame):

input_count = call_frame.contents.args.size
inputs = ctypes.cast(call_frame.contents.args.args, ctypes.POINTER(ctypes.POINTER(XLA_FFI_Buffer)))
inputs = [FfiBuffer(inputs[i].contents) for i in range(input_count)]
inputs = [FfiBuffer(inputs[i].contents, platform=platform) for i in range(input_count)]

output_count = call_frame.contents.rets.size
outputs = ctypes.cast(call_frame.contents.rets.rets, ctypes.POINTER(ctypes.POINTER(XLA_FFI_Buffer)))
outputs = [FfiBuffer(outputs[i].contents) for i in range(output_count)]
outputs = [FfiBuffer(outputs[i].contents, platform=platform) for i in range(output_count)]

ctx = ExecutionContext(call_frame.contents)
ctx = ExecutionContext(call_frame.contents, platform=platform)

func(inputs, outputs, attrs, ctx)

Expand All @@ -1737,12 +1796,26 @@ def ffi_callback(call_frame):
return None

FFI_CCALLFUNC = ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.POINTER(XLA_FFI_CallFrame))
callback_func = FFI_CCALLFUNC(ffi_callback)

def make_callback(platform):
def callback(call_frame):
return ffi_callback(call_frame, platform)

return callback

callback_func_cuda = FFI_CCALLFUNC(make_callback("CUDA"))
with _FFI_REGISTRY_LOCK:
_FFI_CALLBACK_REGISTRY[name] = callback_func
ffi_ccall_address = ctypes.cast(callback_func, ctypes.c_void_p)
ffi_capsule = jax.ffi.pycapsule(ffi_ccall_address.value)
jax.ffi.register_ffi_target(name, ffi_capsule, platform="CUDA")
_FFI_CALLBACK_REGISTRY[name + "_cuda"] = callback_func_cuda
ffi_capsule_cuda = jax.ffi.pycapsule(ctypes.cast(callback_func_cuda, ctypes.c_void_p).value)
jax.ffi.register_ffi_target(name, ffi_capsule_cuda, platform="CUDA")
try:
callback_func_host = FFI_CCALLFUNC(make_callback("Host"))
with _FFI_REGISTRY_LOCK:
_FFI_CALLBACK_REGISTRY[name + "_host"] = callback_func_host
ffi_capsule_host = jax.ffi.pycapsule(ctypes.cast(callback_func_host, ctypes.c_void_p).value)
jax.ffi.register_ffi_target(name, ffi_capsule_host, platform="Host")
except AttributeError:
pass


###############################################################################
Expand Down
23 changes: 20 additions & 3 deletions warp/_src/jax/xla_ffi.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -645,28 +645,45 @@ def jax_dtype_from_ffi(ffi_dtype):
class ExecutionContext:
stage: XLA_FFI_ExecutionStage
stream: int
platform: str

def __init__(self, callframe: XLA_FFI_CallFrame):
def __init__(self, callframe: XLA_FFI_CallFrame, platform="CUDA"):
self.stage = XLA_FFI_ExecutionStage(callframe.stage)
self.stream = get_stream_from_callframe(callframe)
self.platform = platform
Comment thread
greptile-apps[bot] marked this conversation as resolved.
Comment thread
greptile-apps[bot] marked this conversation as resolved.
self.stream = get_stream_from_callframe(callframe) if platform == "CUDA" else None


class FfiBuffer:
dtype: str
data: int
shape: tuple[int]
platform: str

def __init__(self, xla_buffer):
def __init__(self, xla_buffer, platform="CUDA"):
# TODO check if valid
self.dtype = _jnp_dtype_for_xla(xla_buffer.dtype)
self.shape = tuple(xla_buffer.dims[i] for i in range(xla_buffer.rank))
self.data = xla_buffer.data
self.platform = platform

@property
def __cuda_array_interface__(self):
if self.platform != "CUDA":
raise AttributeError("__cuda_array_interface__ is only available on CUDA platform")
return {
"shape": self.shape,
"typestr": self.dtype.char,
"data": (self.data, False),
"version": 2,
}

@property
def __array_interface__(self):
if self.platform != "Host":
raise AttributeError("__array_interface__ is only available on Host platform")
return {
"shape": self.shape,
"typestr": self.dtype.char,
"data": (self.data, False),
"version": 3,
}
Loading
Loading