Description
Add CPU support to wp.jax_kernel() and wp.jax_callable(). A wrapper should work when JAX places a compiled call on CPU and continue to work when the call is placed on CUDA. Users should not need a separate Warp device argument or a different wrapper for each backend.
CPU calls should work inside jax.jit without requiring a CUDA driver or copying buffers to a GPU. Graph modes that do not rely on Warp CUDA graphs should remain available on CPU. CUDA-only graph modes should return a clear error when used with CPU execution.
If an application requests CUDA from a CPU-only Warp build, the call should fail instead of silently running on CPU. Module preload modes should recognize CPU as a supported device.
Context
The JAX FFI wrappers are currently limited to CUDA. CPU-only JAX applications cannot use these APIs inside compiled functions, while applications that run on both CPU and CUDA need backend-specific code.
JAX already chooses where compiled computations run. Supporting CPU execution would let applications move inputs between CPU and CUDA without changing the wrapped function.
Description
Add CPU support to
wp.jax_kernel()andwp.jax_callable(). A wrapper should work when JAX places a compiled call on CPU and continue to work when the call is placed on CUDA. Users should not need a separate Warp device argument or a different wrapper for each backend.CPU calls should work inside
jax.jitwithout requiring a CUDA driver or copying buffers to a GPU. Graph modes that do not rely on Warp CUDA graphs should remain available on CPU. CUDA-only graph modes should return a clear error when used with CPU execution.If an application requests CUDA from a CPU-only Warp build, the call should fail instead of silently running on CPU. Module preload modes should recognize CPU as a supported device.
Context
The JAX FFI wrappers are currently limited to CUDA. CPU-only JAX applications cannot use these APIs inside compiled functions, while applications that run on both CPU and CUDA need backend-specific code.
JAX already chooses where compiled computations run. Supporting CPU execution would let applications move inputs between CPU and CUDA without changing the wrapped function.