From 261734283544531a9fa16c352ae8c137bf44a97b Mon Sep 17 00:00:00 2001 From: Ankit Jain Date: Mon, 6 Jul 2026 21:11:04 +0100 Subject: [PATCH] Add JAX FFI Host (CPU) support 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. (GH-1446) Signed-off-by: Ankit Jain --- CHANGELOG.md | 2 + warp/_src/jax/ffi.py | 191 +++++++++++++------ warp/_src/jax/xla_ffi.py | 23 ++- warp/tests/interop/test_jax.py | 326 +++++++++++++++++++++++++++++++++ 4 files changed, 480 insertions(+), 62 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4dbab6069b..458e09b004 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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 diff --git a/warp/_src/jax/ffi.py b/warp/_src/jax/ffi.py index 2edd8b375b..45f70e4d3c 100644 --- a/warp/_src/jax/ffi.py +++ b/warp/_src/jax/ffi.py @@ -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() @@ -332,7 +344,8 @@ 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: @@ -340,9 +353,10 @@ def __call__(self, *args, output_dims=None, launch_dims=None, vmap_method=None): 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 @@ -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 @@ -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 @@ -448,11 +463,20 @@ 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 @@ -460,20 +484,32 @@ def ffi_callback(self, call_frame): _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()) @@ -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() @@ -715,7 +763,8 @@ 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: @@ -723,9 +772,10 @@ def __call__(self, *args, output_dims=None, vmap_method=None): 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 @@ -733,7 +783,7 @@ def __call__(self, *args, output_dims=None, vmap_method=None): 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 @@ -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 @@ -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") + stream = None # reconstruct the argument list arg_list = [] @@ -906,7 +963,7 @@ def ffi_callback(self, call_frame): # call the Python function with reconstructed arguments with wp.ScopedStream(stream, sync_enter=False): - 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) @@ -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``. """ @@ -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() @@ -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 @@ -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 @@ -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) @@ -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 ############################################################################### diff --git a/warp/_src/jax/xla_ffi.py b/warp/_src/jax/xla_ffi.py index a032c5b973..ad0ec8a902 100644 --- a/warp/_src/jax/xla_ffi.py +++ b/warp/_src/jax/xla_ffi.py @@ -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 + 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, + } diff --git a/warp/tests/interop/test_jax.py b/warp/tests/interop/test_jax.py index 113b0c6787..0a0b6fc232 100644 --- a/warp/tests/interop/test_jax.py +++ b/warp/tests/interop/test_jax.py @@ -590,6 +590,15 @@ def scale_sum_square_kernel_subscript(a: wp.array[float], b: wp.array[float], s: c[tid] = (a[tid] * s + b[tid]) ** 2.0 +TILE_ONES_N = 64 + + +@wp.kernel +def tile_ones_store_kernel(output: wp.array[float]): + t = wp.tile_ones(dtype=float, shape=TILE_ONES_N) + wp.tile_store(output, t) + + # The Python function to call. # Note the argument annotations, just like Warp kernels. def scale_func( @@ -1257,6 +1266,282 @@ def warp_func(inputs, outputs, attrs, ctx): assert_np_equal(d, 2 * np.arange(ARRAY_SIZE, dtype=np.float32).reshape((ARRAY_SIZE // 2, 2))) +@unittest.skipUnless(_jax_version() >= (0, 5, 0), "Jax version too old") +@unittest.skipUnless(wp.is_cpu_available(), "Requires CPU backend") +def test_ffi_jax_kernel_host_add(test, device): + import jax.numpy as jp # noqa: PLC0415 + + jax_kernel = wp.jax_kernel + + jax_add = jax_kernel(add_kernel) + + @jax.jit + def f(): + n = ARRAY_SIZE + a = jp.arange(n, dtype=jp.float32) + b = jp.ones(n, dtype=jp.float32) + return jax_add(a, b) + + with jax.default_device(jax.devices("cpu")[0]): + (y,) = f() + + jax.block_until_ready(y) + + result = np.asarray(y) + expected = np.arange(1, ARRAY_SIZE + 1, dtype=np.float32) + + assert_np_equal(result, expected) + + +@unittest.skipUnless(_jax_version() >= (0, 5, 0), "Jax version too old") +@unittest.skipUnless(wp.is_cpu_available(), "Requires CPU backend") +def test_ffi_jax_kernel_host_sincos(test, device): + import jax.numpy as jp # noqa: PLC0415 + + jax_kernel = wp.jax_kernel + + jax_sincos = jax_kernel(sincos_kernel, num_outputs=2) + + @jax.jit + def f(): + n = ARRAY_SIZE + # bounded inputs to avoid precision issues with float32 for large values + a = jp.linspace(-jp.pi, jp.pi, n, dtype=jp.float32) + return jax_sincos(a) + + with jax.default_device(jax.devices("cpu")[0]): + s, c = f() + + jax.block_until_ready([s, c]) + + ref_a = np.linspace(-np.pi, np.pi, ARRAY_SIZE, dtype=np.float32) + ref_s = np.sin(ref_a) + ref_c = np.cos(ref_a) + + assert_np_equal(np.asarray(s), ref_s, tol=1e-4) + assert_np_equal(np.asarray(c), ref_c, tol=1e-4) + + +@unittest.skipUnless(_jax_version() >= (0, 5, 0), "Jax version too old") +@unittest.skipUnless(wp.is_cpu_available(), "Requires CPU backend") +def test_ffi_jax_kernel_host_in_out(test, device): + import jax.numpy as jp # noqa: PLC0415 + + jax_kernel = wp.jax_kernel + + # b is in-out, c is out only + jax_in_out = jax_kernel(in_out_kernel, num_outputs=2, in_out_argnames=["b"]) + + @jax.jit + def f(): + n = ARRAY_SIZE + a = jp.ones(n, dtype=jp.float32) + b = jp.arange(n, dtype=jp.float32) + return jax_in_out(a, b) + + with jax.default_device(jax.devices("cpu")[0]): + b, c = f() + + jax.block_until_ready([b, c]) + + assert_np_equal(np.asarray(b), np.arange(1, ARRAY_SIZE + 1, dtype=np.float32)) + assert_np_equal(np.asarray(c), np.full(ARRAY_SIZE, 2, dtype=np.float32)) + + +@unittest.skipUnless(_jax_version() >= (0, 5, 0), "Jax version too old") +@unittest.skipUnless(wp.is_cpu_available(), "Requires CPU backend") +def test_ffi_jax_kernel_host_scale_vec_constant(test, device): + import jax.numpy as jp # noqa: PLC0415 + + jax_kernel = wp.jax_kernel + + jax_scale_vec = jax_kernel(scale_vec_kernel) + + @jax.jit + def f(): + a = jp.arange(ARRAY_SIZE, dtype=jp.float32).reshape((ARRAY_SIZE // 2, 2)) + s = 2.0 + return jax_scale_vec(a, s) + + with jax.default_device(jax.devices("cpu")[0]): + (b,) = f() + + jax.block_until_ready(b) + + expected = 2 * np.arange(ARRAY_SIZE, dtype=np.float32).reshape((ARRAY_SIZE // 2, 2)) + + assert_np_equal(b, expected) + + +@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_scale_constant(test, device): + import jax.numpy as jp # noqa: PLC0415 + + jax_callable = wp.jax_callable + + jax_func = jax_callable(scale_func, num_outputs=2) + + @jax.jit + def f(): + a = jp.arange(ARRAY_SIZE, dtype=jp.float32) + b = jp.arange(ARRAY_SIZE, dtype=jp.float32).reshape((ARRAY_SIZE // 2, 2)) # array of vec2 + s = 2.0 + + # output shapes + output_dims = {"c": a.shape, "d": b.shape} + + return jax_func(a, b, s, output_dims=output_dims) + + with jax.default_device(jax.devices("cpu")[0]): + result1, result2 = f() + + jax.block_until_ready([result1, result2]) + + assert_np_equal(result1, 2 * np.arange(ARRAY_SIZE, dtype=np.float32)) + assert_np_equal(result2, 2 * np.arange(ARRAY_SIZE, dtype=np.float32).reshape((ARRAY_SIZE // 2, 2))) + + +@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_in_out(test, device): + import jax.numpy as jp # noqa: PLC0415 + + jax_callable = wp.jax_callable + + # alpha is static (scalar attribute), b is in-out, c is output only + jax_func = jax_callable(in_out_func, num_outputs=2, in_out_argnames=["b"]) + + @jax.jit + def f(): + n = ARRAY_SIZE + a = jp.ones(n, dtype=jp.float32) + b = jp.arange(n, dtype=jp.float32) + return jax_func(a, b) + + with jax.default_device(jax.devices("cpu")[0]): + b, c = f() + + jax.block_until_ready([b, c]) + + assert_np_equal(np.asarray(b), np.arange(1, ARRAY_SIZE + 1, dtype=np.float32)) + assert_np_equal(np.asarray(c), np.full(ARRAY_SIZE, 2, dtype=np.float32)) + + +@unittest.skipUnless(_jax_version() >= (0, 5, 0), "Jax version too old") +@unittest.skipUnless(wp.is_cpu_available(), "Requires CPU backend") +def test_ffi_callback_host(test, device): + jax = _import_jax() + import jax.numpy as jp # noqa: PLC0415 + + from warp._src.jax.ffi import register_ffi_callback # noqa: PLC0415 + + # the Python function to call + def warp_func(inputs, outputs, attrs, ctx): + # input arrays + a = inputs[0] + b = inputs[1] + + test.assertEqual(ctx.stream, None) + test.assertTrue(hasattr(a, "__array_interface__")) + test.assertFalse(hasattr(a, "__cuda_array_interface__")) + + # scalar attributes + s = attrs["scale"] + + # output arrays + c = outputs[0] + d = outputs[1] + + # launch with arrays of scalars + wp.launch(scale_kernel, dim=a.shape, inputs=[a, s], outputs=[c], device="cpu") + + # launch with arrays of vec2 + # NOTE: the input shapes are from JAX arrays, we need to strip the inner dimension for vec2 arrays + wp.launch(scale_vec_kernel, dim=b.shape[0], inputs=[b, s], outputs=[d], device="cpu") + + # register callback + register_ffi_callback("warp_func_host", warp_func) + + n = ARRAY_SIZE + + with jax.default_device(jax.devices("cpu")[0]): + # inputs + a = jp.arange(n, dtype=jp.float32) + b = jp.arange(n, dtype=jp.float32).reshape((n // 2, 2)) # array of wp.vec2 + s = 2.0 + + # set up call + out_types = [ + jax.ShapeDtypeStruct(a.shape, jp.float32), + jax.ShapeDtypeStruct(b.shape, jp.float32), # array of wp.vec2 + ] + call = jax.ffi.ffi_call("warp_func_host", out_types) + + # call it + c, d = call(a, b, scale=s) + + jax.block_until_ready([c, d]) + + assert_np_equal(np.asarray(c), 2 * np.arange(ARRAY_SIZE, dtype=np.float32)) + assert_np_equal(np.asarray(d), 2 * np.arange(ARRAY_SIZE, dtype=np.float32).reshape((ARRAY_SIZE // 2, 2))) + + +@unittest.skipUnless(_jax_version() >= (0, 5, 0), "Jax version too old") +@unittest.skipUnless(wp.is_cpu_available(), "Requires CPU backend") +def test_ffi_jax_kernel_host_tile_ones(test, device): + """Regression: Host callbacks must compile WP_TILE_BLOCK_DIM=1 so that + explicitly shaped tile operations (e.g. tile_ones(shape=64)) write all + elements, not just lane 0's portion of the tile.""" + + jax_kernel = wp.jax_kernel + + jax_tile_ones = jax_kernel(tile_ones_store_kernel, launch_dims=1) + + @jax.jit + def f(): + return jax_tile_ones(output_dims=(TILE_ONES_N,)) + + with jax.default_device(jax.devices("cpu")[0]): + (y,) = f() + + jax.block_until_ready(y) + + result = np.asarray(y) + expected = np.ones(TILE_ONES_N, dtype=np.float32) + + assert_np_equal(result, expected) + + +@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_tile_ones(test, device): + """Regression: Host callbacks must compile WP_TILE_BLOCK_DIM=1 so that + explicitly shaped tile operations (e.g. tile_ones(shape=64)) write all + elements, not just lane 0's portion of the tile.""" + + jax_callable = wp.jax_callable + + def tile_ones_store_func(output: wp.array[float]): + wp.launch_tiled(tile_ones_store_kernel, dim=1, block_dim=1, inputs=[output], device="cpu") + + jax_tile_ones = jax_callable(tile_ones_store_func, graph_mode=wp.JaxCallableGraphMode.JAX) + + @jax.jit + def f(): + return jax_tile_ones(output_dims=(TILE_ONES_N,)) + + with jax.default_device(jax.devices("cpu")[0]): + (y,) = f() + + jax.block_until_ready(y) + + result = np.asarray(y) + expected = np.ones(TILE_ONES_N, dtype=np.float32) + + assert_np_equal(result, expected) + + @unittest.skipUnless(_jax_version() >= (0, 5, 0), "Jax version too old") def test_ffi_jax_kernel_autodiff_simple(test, device): if device.ordinal > 0: @@ -2529,6 +2814,47 @@ class TestJax(unittest.TestCase): devices=jax_compatible_cuda_devices, ) + # ffi Host (CPU) tests — always registered (not gated on CUDA availability) + add_function_test(TestJax, "test_ffi_jax_kernel_host_add", test_ffi_jax_kernel_host_add, devices=None) + add_function_test(TestJax, "test_ffi_jax_kernel_host_sincos", test_ffi_jax_kernel_host_sincos, devices=None) + add_function_test(TestJax, "test_ffi_jax_kernel_host_in_out", test_ffi_jax_kernel_host_in_out, devices=None) + add_function_test( + TestJax, + "test_ffi_jax_kernel_host_scale_vec_constant", + test_ffi_jax_kernel_host_scale_vec_constant, + devices=None, + ) + add_function_test( + TestJax, + "test_ffi_jax_callable_host_scale_constant", + test_ffi_jax_callable_host_scale_constant, + devices=None, + ) + add_function_test( + TestJax, + "test_ffi_jax_callable_host_in_out", + test_ffi_jax_callable_host_in_out, + devices=None, + ) + add_function_test( + TestJax, + "test_ffi_callback_host", + test_ffi_callback_host, + devices=None, + ) + add_function_test( + TestJax, + "test_ffi_jax_kernel_host_tile_ones", + test_ffi_jax_kernel_host_tile_ones, + devices=None, + ) + add_function_test( + TestJax, + "test_ffi_jax_callable_host_tile_ones", + test_ffi_jax_callable_host_tile_ones, + devices=None, + ) + # bfloat16 tests require arch >= 80 bf16_jax_devices = [d for d in jax_compatible_devices if d.is_cpu or (d.is_cuda and d.arch >= 80)] if bf16_jax_devices: