Add JAX FFI Host support#1446
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThis pull request extends JAX experimental FFI callbacks to support both CUDA and Host (CPU) execution paths. Callbacks now accept a ChangesCUDA and Host FFI Execution
🎯 4 (Complex) | ⏱️ ~45 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@warp/_src/jax_experimental/ffi.py`:
- Around line 1772-1777: In register_ffi_callback, the Host ffi capsule is
constructed from the wrong variable (ffi_capsule_host) causing a NameError;
change the construction to use the host ccall address value by calling
jax.ffi.pycapsule(ffi_ccall_address_host.value) and then register that capsule
with jax.ffi.register_ffi_target(name, ffi_capsule_host, platform="Host") so the
Host path mirrors the CUDA path (refer to ffi_ccall_address_host,
ffi_capsule_host, register_ffi_target).
- Around line 629-632: The code assigns ffi_ccall_address_host then creates
ffi_capsule_host but mistakenly uses ffi_capsule_host.value (self-referential
NameError); change the capsule creation to use the previously computed address
value (ffi_ccall_address_host.value) so the lines around callback_func_host,
ffi_ccall_address_host, ffi_capsule_host and the
jax.ffi.register_ffi_target(self.name, ffi_capsule_host, platform="Host") call
use ffi_ccall_address_host.value when constructing the pycapsule for the Host
callback that wraps FFI_CCALLFUNC and calls self.ffi_callback.
- Around line 226-229: The Host FFI registration references an undefined
variable: replace the erroneous creation of ffi_capsule_host (currently using
ffi_capsule_host.value) with a capsule built from the c_void_p address you just
made; specifically, after creating callback_func_host and
ffi_ccall_address_host, set ffi_capsule_host =
jax.ffi.pycapsule(ffi_ccall_address_host.value) and then call
jax.ffi.register_ffi_target(self.name, ffi_capsule_host, platform="Host") so the
capsule is created from the correct address (symbols: callback_func_host,
FFI_CCALLFUNC, ffi_ccall_address_host, ffi_capsule_host,
jax.ffi.register_ffi_target, self.name, self.ffi_callback).
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yml
Review profile: CHILL
Plan: Enterprise
Run ID: d2c93d0f-cbef-4891-b361-773ad8140f6c
📒 Files selected for processing (1)
warp/_src/jax_experimental/ffi.py
|
pre-commit.ci autofix |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@warp/tests/interop/test_jax.py`:
- Around line 2422-2450: The host-only test registrations for TestJax (the
add_function_test calls registering test_ffi_jax_kernel_host_add,
test_ffi_jax_kernel_host_sincos, test_ffi_jax_kernel_host_in_out,
test_ffi_jax_kernel_host_scale_vec_constant,
test_ffi_jax_callable_host_scale_constant, and
test_ffi_jax_callable_host_in_out) are currently inside the
jax_compatible_cuda_devices conditional; move these specific add_function_test
calls out of that CUDA-only if block so they are always registered on CPU-only
setups, keeping the existing device=None argument and leaving CUDA/GPU-specific
registrations inside the original conditional.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yml
Review profile: CHILL
Plan: Enterprise
Run ID: fcfb2e10-e509-41b9-b675-c8960056b6be
📒 Files selected for processing (1)
warp/tests/interop/test_jax.py
|
pre-commit.ci autofix |
|
Please add a |
|
Please squash this PR down to a single coherent commit before merge. |
|
Please rebase onto current |
nvlukasz
left a comment
nvlukasz
left a comment
There was a problem hiding this comment.
Thank you for the contribution. Please address outstanding comments or let us know if you are unable to do so.
@nvlukasz , please allow till the end of this week to address the outstanding comments. I apologise for the delay. |
loney7
force-pushed
the
loney7/ffi-host-support
branch
2 times, most recently
from
766a819 to
cc5aa99
Compare
|
pre-commit.ci autofix |
loney7
force-pushed
the
loney7/ffi-host-support
branch
from
cc5aa99 to
c06e8d4
Compare
|
pre-commit.ci autofix |
loney7
force-pushed
the
loney7/ffi-host-support
branch
from
420290b to
38f4e20
Compare
|
pre-commit.ci autofix |
|
Hi @nvlukasz wondering if you could take another pass :) |
loney7
force-pushed
the
loney7/ffi-host-support
branch
from
df57595 to
f752ab9
Compare
|
pre-commit.ci autofix |
1 similar comment
|
pre-commit.ci autofix |
|
Want your agent to iterate on Greptile's feedback? Try greploops. |
loney7
force-pushed
the
loney7/ffi-host-support
branch
from
5d52c5c to
865d5ef
Compare
|
/ok to test 865d5ef |
loney7
force-pushed
the
loney7/ffi-host-support
branch
from
d6973f2 to
62590b8
Compare
|
Could you also fix this Host Minimal repro: import jax
import jax.numpy as jnp
import numpy as np
import warp as wp
@wp.kernel
def add_kernel(
a: wp.array[wp.float32],
b: wp.array[wp.float32],
out: wp.array[wp.float32],
):
tid = wp.tid()
out[tid] = a[tid] + b[tid]
jax_add = wp.jax_kernel(add_kernel)
@jax.jit
def f(a, b):
(out,) = jax_add(a, b)
return out
cpu = jax.devices("cpu")[0]
# Leave the effective JAX default device as CUDA, but place inputs on CPU.
a = jax.device_put(jnp.arange(16, dtype=jnp.float32), cpu)
b = jax.device_put(jnp.ones(16, dtype=jnp.float32), cpu)
out = f(a, b)
jax.block_until_ready(out)
np.testing.assert_allclose(np.asarray(out), np.arange(1, 17, dtype=np.float32))On this branch, The Host callback should ensure the kernel module is loaded for CPU before calling |
|
Could you also fix the generic Host FFI callback buffer path? The Host target is registered, but the buffers passed to the user callback are still Minimal repro: import warnings
import jax
import jax.numpy as jnp
import warp as wp
with warnings.catch_warnings():
warnings.simplefilter("ignore", DeprecationWarning)
from warp.jax_experimental import register_ffi_callback
@wp.kernel
def double_kernel(inp: wp.array[wp.float32], out: wp.array[wp.float32]):
tid = wp.tid()
out[tid] = 2.0 * inp[tid]
def warp_callback(inputs, outputs, attrs, ctx):
inp = inputs[0]
out = outputs[0]
assert ctx.stream is None
assert hasattr(inp, "__cuda_array_interface__")
assert not hasattr(inp, "__array_interface__")
wp.launch(double_kernel, dim=inp.shape[0], inputs=[inp], outputs=[out], device="cpu")
register_ffi_callback("warp_generic_host_buffer_repro", warp_callback)
@jax.jit
def f(x):
out_type = jax.ShapeDtypeStruct(x.shape, x.dtype)
call = jax.ffi.ffi_call("warp_generic_host_buffer_repro", out_type)
return call(x)
with jax.default_device(jax.devices("cpu")[0]):
x = jnp.arange(16, dtype=jnp.float32)
out = f(x)
jax.block_until_ready(out)On this branch, the Host callback is reached, but the CPU launch fails while packing the Can the Host generic callback path expose CPU-compatible buffers, e.g. via |
Add support for running JAX FFI callbacks on the CPU (Host) platform in addition to CUDA. Changes: - Register FFI targets for both "CUDA" and "Host" platforms in FfiKernel, FfiCallable, and register_ffi_callback. - Handle the "Host" platform case in ffi_callback by using the CPU device and bypassing CUDA-specific features like streams and graphs. - Update FfiCallable to reconstruct arguments and execute the function on the CPU when running on the Host platform. - Make FfiBuffer platform-aware: expose __array_interface__ for Host and __cuda_array_interface__ for CUDA. - Ensure kernel module is loaded on the target device before getting hooks, fixing the case where JAX default device is CUDA but the callback runs on Host. - Guard COMMAND_BUFFER_COMPATIBLE traits and graph_mode setup behind platform == "CUDA" checks. - Add CPU/Host FFI tests covering kernels, callables, and callbacks. (NVIDIAGH-1446) Signed-off-by: Ankit Jain <kitsrish@google.com>
loney7
force-pushed
the
loney7/ffi-host-support
branch
from
62590b8 to
c9d6224
Compare
loney7
force-pushed
the
loney7/ffi-host-support
branch
from
c9d6224 to
fbca388
Compare
Add support for running JAX FFI callbacks on the CPU (Host) platform in addition to CUDA. Changes: - Register FFI targets for both "CUDA" and "Host" platforms in FfiKernel, FfiCallable, and register_ffi_callback. - Handle the "Host" platform case in ffi_callback by using the CPU device and bypassing CUDA-specific features like streams and graphs. - Update FfiCallable to reconstruct arguments and execute the function on the CPU when running on the Host platform. - Make FfiBuffer platform-aware: expose __array_interface__ for Host and __cuda_array_interface__ for CUDA. - Ensure kernel module is loaded on the target device before getting hooks, fixing the case where JAX default device is CUDA but the callback runs on Host. - Guard COMMAND_BUFFER_COMPATIBLE traits and graph_mode setup behind platform == "CUDA" checks. - Add CPU/Host FFI tests covering kernels, callables, and callbacks. (NVIDIAGH-1446) Signed-off-by: Ankit Jain <kitsrish@google.com>
|
@shi-eric , can you please have a look |
|
Could you hard-code
Something along these lines: 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)Could you also add a Host regression test that stores an explicitly shaped tile, such as |
loney7
force-pushed
the
loney7/ffi-host-support
branch
from
fbca388 to
7253c19
Compare
Add support for running JAX FFI callbacks on the CPU (Host) platform in addition to CUDA. Changes: - Register FFI targets for both "CUDA" and "Host" platforms in FfiKernel, FfiCallable, and register_ffi_callback. - Handle the "Host" platform case in ffi_callback by using the CPU device and bypassing CUDA-specific features like streams and graphs. - Update FfiCallable to reconstruct arguments and execute the function on the CPU when running on the Host platform. - Make FfiBuffer platform-aware: expose __array_interface__ for Host and __cuda_array_interface__ for CUDA. - Ensure kernel module is loaded on the target device before getting hooks, fixing the case where JAX default device is CUDA but the callback runs on Host. - Guard COMMAND_BUFFER_COMPATIBLE traits and graph_mode setup behind platform == "CUDA" checks. - Add CPU/Host FFI tests covering kernels, callables, and callbacks. (NVIDIAGH-1446) Signed-off-by: Ankit Jain <kitsrish@google.com>
loney7
force-pushed
the
loney7/ffi-host-support
branch
from
7253c19 to
caaffcd
Compare
Add support for running JAX FFI callbacks on the CPU (Host) platform in addition to CUDA. Changes: - Register FFI targets for both "CUDA" and "Host" platforms in FfiKernel, FfiCallable, and register_ffi_callback. - Handle the "Host" platform case in ffi_callback by using the CPU device and bypassing CUDA-specific features like streams and graphs. - Update FfiCallable to reconstruct arguments and execute the function on the CPU when running on the Host platform. - Make FfiBuffer platform-aware: expose __array_interface__ for Host and __cuda_array_interface__ for CUDA. - Ensure kernel module is loaded on the target device before getting hooks, fixing the case where JAX default device is CUDA but the callback runs on Host. - Guard COMMAND_BUFFER_COMPATIBLE traits and graph_mode setup behind platform == "CUDA" checks. - Add CPU/Host FFI tests covering kernels, callables, and callbacks. (NVIDIAGH-1446) Signed-off-by: Ankit Jain <kitsrish@google.com>
Add support for running JAX FFI callbacks on the CPU (Host) platform in addition to CUDA. Changes: - Register FFI targets for both "CUDA" and "Host" platforms in FfiKernel, FfiCallable, and register_ffi_callback. - Handle the "Host" platform case in ffi_callback by using the CPU device and bypassing CUDA-specific features like streams and graphs. - Update FfiCallable to reconstruct arguments and execute the function on the CPU when running on the Host platform. - Make FfiBuffer platform-aware: expose __array_interface__ for Host and __cuda_array_interface__ for CUDA. - Ensure kernel module is loaded on the target device before getting hooks, fixing the case where JAX default device is CUDA but the callback runs on Host. - Guard COMMAND_BUFFER_COMPATIBLE traits and graph_mode setup behind platform == "CUDA" checks. - Add CPU/Host FFI tests covering kernels, callables, and callbacks. (NVIDIAGH-1446) Signed-off-by: Ankit Jain <kitsrish@google.com>
loney7
force-pushed
the
loney7/ffi-host-support
branch
from
caaffcd to
47b367c
Compare
loney7
force-pushed
the
loney7/ffi-host-support
branch
2 times, most recently
from
ca5e33a to
90b7047
Compare
Add support for running JAX FFI callbacks on the CPU (Host) platform in addition to CUDA. Changes: - Register FFI targets for both "CUDA" and "Host" platforms in FfiKernel, FfiCallable, and register_ffi_callback. - Handle the "Host" platform case in ffi_callback by using the CPU device and bypassing CUDA-specific features like streams and graphs. - Update FfiCallable to reconstruct arguments and execute the function on the CPU when running on the Host platform. - Make FfiBuffer platform-aware: expose __array_interface__ for Host and __cuda_array_interface__ for CUDA. - Ensure kernel module is loaded on the target device before getting hooks, fixing the case where JAX default device is CUDA but the callback runs on Host. - Guard COMMAND_BUFFER_COMPATIBLE traits and graph_mode setup behind platform == "CUDA" checks. - Add CPU/Host FFI tests covering kernels, callables, and callbacks. (NVIDIAGH-1446) Signed-off-by: Ankit Jain <kitsrish@google.com>
Add support for running JAX FFI callbacks on the CPU (Host) platform in addition to CUDA. Changes: - Register FFI targets for both "CUDA" and "Host" platforms in FfiKernel, FfiCallable, and register_ffi_callback. - Handle the "Host" platform case in ffi_callback by using the CPU device and bypassing CUDA-specific features like streams and graphs. - Update FfiCallable to reconstruct arguments and execute the function on the CPU when running on the Host platform. - Make FfiBuffer platform-aware: expose __array_interface__ for Host and __cuda_array_interface__ for CUDA. - Ensure kernel module is loaded on the target device before getting hooks, fixing the case where JAX default device is CUDA but the callback runs on Host. - Guard COMMAND_BUFFER_COMPATIBLE traits and graph_mode setup behind platform == "CUDA" checks. - Add CPU/Host FFI tests covering kernels, callables, and callbacks. (NVIDIAGH-1446) Signed-off-by: Ankit Jain <kitsrish@google.com>
loney7
force-pushed
the
loney7/ffi-host-support
branch
from
90b7047 to
74bd896
Compare
Add support for running JAX FFI callbacks on the CPU (Host) platform in addition to CUDA. Changes: - Register FFI targets for both "CUDA" and "Host" platforms in FfiKernel, FfiCallable, and register_ffi_callback. - Handle the "Host" platform case in ffi_callback by using the CPU device and bypassing CUDA-specific features like streams and graphs. - Update FfiCallable to reconstruct arguments and execute the function on the CPU when running on the Host platform. - Make FfiBuffer platform-aware: expose __array_interface__ for Host and __cuda_array_interface__ for CUDA. - Ensure kernel module is loaded on the target device before getting hooks, fixing the case where JAX default device is CUDA but the callback runs on Host. - Guard COMMAND_BUFFER_COMPATIBLE traits and graph_mode setup behind platform == "CUDA" checks. - Add CPU/Host FFI tests covering kernels, callables, and callbacks. (NVIDIAGH-1446) Signed-off-by: Ankit Jain <kitsrish@google.com>
loney7
force-pushed
the
loney7/ffi-host-support
branch
from
74bd896 to
2617342
Compare
|
/ok to test 2617342 |
| @@ -906,7 +963,7 @@ def ffi_callback(self, call_frame): | |||
|
|
|||
| # call the Python function with reconstructed arguments | |||
| with wp.ScopedStream(stream, sync_enter=False): | |||
There was a problem hiding this comment.
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.
| 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)
4 participants
Description
This PR adds support for running JAX FFI callbacks on the CPU (Host) in addition to CUDA.
Changes:
"CUDA"and"Host"platforms inregister_ffi_callback."Host"platform case inffi_callbackby using the CPU device and bypassing CUDA-specific features like streams and graphs.FfiCallableto reconstruct arguments and execute the function on the CPU when running on the Host platform.Checklist
Test plan
You can verify these changes by running the JAX interop tests which include FFI tests. Ensure they pass on both CPU and GPU (if available).