Add CPU support for JAX FFI (GH-1661)#1662
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:
📝 WalkthroughWalkthroughJAX FFI wrappers now support CPU and CUDA execution selected by JAX, with platform-specific registration, module preloading, graph-mode validation, expanded interoperability tests, updated documentation, and CI JAX requirements. ChangesJAX CPU FFI support
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant JAX
participant Warp FFI
participant Warp Runtime
participant CPU or CUDA Device
JAX->>Warp FFI: Lower and invoke jax_kernel or jax_callable
Warp FFI->>Warp Runtime: Select platform and load module
Warp Runtime->>CPU or CUDA Device: Execute kernel or callable
CPU or CUDA Device-->>JAX: Return output buffers
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
warp/tests/interop/test_jax.py (1)
131-131: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse regular unittest methods for tests without harness-selected devices.
The preload and CPU-subprocess tests ignore the injected
device; make themTestJaxmethods rather than registering them withadd_function_test().As per coding guidelines, use standard
unittest.TestCasemethods for fixed-device tests and reserveadd_function_test()for tests running across multiple devices viaget_test_devices().Also applies to: 184-184, 1290-1292, 2462-2475
🤖 Prompt for 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. In `@warp/tests/interop/test_jax.py` at line 131, Convert the preload and CPU-subprocess tests, including test_ffi_module_preload_modes and the additionally referenced tests, into methods on TestJax instead of registering them through add_function_test(). Remove their harness registration and unused injected device parameters, while preserving standard unittest behavior for these fixed-device tests.Source: Coding guidelines
🤖 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 1279-1285: Update the auxiliary subprocess invocation in the test
around subprocess.run to include an appropriate timeout value, ensuring JAX
initialization or compilation hangs terminate before the CI job timeout while
preserving the existing output capture and error-handling behavior.
---
Nitpick comments:
In `@warp/tests/interop/test_jax.py`:
- Line 131: Convert the preload and CPU-subprocess tests, including
test_ffi_module_preload_modes and the additionally referenced tests, into
methods on TestJax instead of registering them through add_function_test().
Remove their harness registration and unused injected device parameters, while
preserving standard unittest behavior for these fixed-device tests.
🪄 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: 1c570803-f266-453c-809f-2812af04ea44
📒 Files selected for processing (7)
.github/workflows/ci.yml.gitlab-ci.ymlCHANGELOG.mddocs/user_guide/interoperability_jax.rstwarp/_src/jax/ffi.pywarp/tests/interop/aux_test_jax_cpu_ffi.pywarp/tests/interop/test_jax.py
Greptile SummaryThis PR adds CPU execution support for
Confidence Score: 5/5Safe to merge — changes are well-scoped, thoroughly tested, and preserve all existing CUDA behavior. Platform dispatch is clean: ctypes closures capture platform via default-argument (no late-binding bug), _load_ffi_module correctly derives block_dim from device type, and CPU early-return paths in FfiKernel and FfiCallable are symmetric with CUDA paths. Refactored _build_arg_list matches original inline code exactly. Module deduplication uses id(device) which is stable since Warp devices are singletons. Test coverage spans unit mocks, in-process CPU/CUDA dispatch, subprocess CPU isolation, and graph-mode boundary validation. No files require special attention. Important Files Changed
Reviews (9): Last reviewed commit: "Add CPU support for JAX FFI (GH-1661)" | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
warp/tests/interop/test_jax.py (1)
2365-2382: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAssert that the unavailable CUDA backend is actually exercised.
The expected loaded-device list remains
[warp_cpu]even if a refactor stops probing CUDA, so this does not verify the intended CUDA-initialization fallback. Record backend requests and assert that"cuda"was requested.Proposed test strengthening
+ requested_backends = [] + class FakeJax: `@staticmethod` def local_devices(process_index=None, backend=None, host_id=None): + requested_backends.append(backend) if backend == "cpu": return [cpu_0, cpu_1] if backend == "cuda": raise RuntimeError("CUDA backend is unavailable") @@ ffi_module._preload_ffi_module(module, wp.JaxModulePreloadMode.ALL_DEVICES) self.assertEqual(module.loaded_devices, [warp_cpu]) + self.assertIn("cuda", requested_backends)🤖 Prompt for 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. In `@warp/tests/interop/test_jax.py` around lines 2365 - 2382, Strengthen the all-devices test around FakeJax.local_devices by recording each requested backend, then assert that "cuda" was requested in addition to preserving the existing loaded_devices assertion. Keep the CUDA RuntimeError path exercised so the test verifies fallback behavior rather than only the final CPU device list.
🤖 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.
Nitpick comments:
In `@warp/tests/interop/test_jax.py`:
- Around line 2365-2382: Strengthen the all-devices test around
FakeJax.local_devices by recording each requested backend, then assert that
"cuda" was requested in addition to preserving the existing loaded_devices
assertion. Keep the CUDA RuntimeError path exercised so the test verifies
fallback behavior rather than only the final CPU device list.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yml
Review profile: CHILL
Plan: Enterprise
Run ID: 7e103a80-5c3e-457d-ab99-2e9ec0ef7577
📒 Files selected for processing (2)
warp/_src/jax/ffi.pywarp/tests/interop/test_jax.py
🚧 Files skipped from review as they are similar to previous changes (1)
- warp/_src/jax/ffi.py
335a120 to
cead8ea
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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/ffi.py`:
- Line 166: Update the preloading calls around module.load in the relevant FFI
initialization path to use _load_ffi_module() instead, including the additional
calls at the referenced locations. Preserve the existing device and
block-dimension arguments so cached module build failures propagate rather than
being silently ignored.
In `@warp/tests/interop/test_jax.py`:
- Line 1108: Update every skip message associated with the visible
_jax_version() decorators to capitalize the proper name, replacing “Jax version
too old” with “JAX version too old” consistently across all occurrences.
🪄 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: 70802b3e-4b65-4b92-a19a-c6105fb24285
📒 Files selected for processing (7)
.github/workflows/ci.yml.gitlab-ci.ymlCHANGELOG.mddocs/user_guide/interoperability_jax.rstwarp/_src/jax/ffi.pywarp/tests/interop/aux_test_jax_cpu_ffi.pywarp/tests/interop/test_jax.py
🚧 Files skipped from review as they are similar to previous changes (5)
- warp/tests/interop/aux_test_jax_cpu_ffi.py
- CHANGELOG.md
- .github/workflows/ci.yml
- .gitlab-ci.yml
- docs/user_guide/interoperability_jax.rst
a4c54a4 to
726ba2f
Compare
JAX FFI wrappers previously assumed CUDA execution, blocking CPU-only applications and forcing mixed-device programs onto backend-specific paths. Dispatch callbacks using the device selected by JAX, compile CPU kernels with a single-lane block dimension, and preserve CUDA graph caching and module preloading behavior. Report module and backend failures through clear FFI errors. Cover CPU, mixed-device, graph replay, module loading, and CI paths, and document the supported behavior. Signed-off-by: Eric Shi <ershi@nvidia.com>
726ba2f to
cdd26d6
Compare
1 participant
Description
Add CPU execution support for
wp.jax_kernel()andwp.jax_callable()through JAX FFI. The same wrapper is registered for CPU and CUDA, allowing
JAX to select the implementation from the device used during lowering.
Closes #1661.
Related work: #1446 covers the same feature area. This implementation was
developed independently. Its CPU shaped-tile correction follows the
block_dim=1issue reported while reviewing that PR:#1446 (comment)
Changes
buffers, while preserving existing CUDA behavior.
single-lane block dimension so shaped tile operations cover their full
extent.
clearly when Warp lacks CUDA support.
preloading-disabled, and shaped-tile user workflows.
Unreleased changelog entry.
Checklist
Unreleasedsection.Validation summary
146 tests passed, with 12 expected skips.
preloading disabled, plus the CPU shaped-tile, CPU callable graph-mode, and
mixed CPU/CUDA regression tests. All four passed.
passed.
Bug fix
Without the CPU-specific block dimension, shaped tile stores only populate the
first element:
New feature / enhancement
The same JAX FFI wrapper can now execute on CPU or CUDA according to the input
device:
Summary by CodeRabbit
wp.jax_kernel()andwp.jax_callable(), with automatic dispatch based on the JAX device.wp.jax_kernel()andwp.jax_callable().