[envpool] migrate XLA custom calls to OpenXLA FFI - #343
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e6fd862740
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| handle_value = int.from_bytes(handle, byteorder=sys.byteorder, signed=False) | ||
|
|
||
| def call(*args: Any) -> Any: | ||
| return xla_func(*args) | ||
| return xla_func(*args, handle=handle_value) |
There was a problem hiding this comment.
ffi_call() now threads the env handle through a named attribute, and CustomCall::ResolveHandle() decodes that attribute with attrs.get<std::int64_t>("handle"). Here we pass handle_value as a plain Python int; on the default JAX configuration (jax_enable_x64=False), Python ints are 32-bit, so the typed-FFI attribute is either rejected as the wrong type or truncated before the pointer cast on 64-bit hosts. In that configuration every send/recv/step call will fail unless users globally enable x64. Use an explicit np.int64/np.uint64 scalar for the attribute.
Useful? React with 👍 / 👎.
1 participant
Summary
jaxlibwheel repos.openxla/xlasource snapshot thatjax-v0.9.2already references.pip_requirements_*_jaxlibheaders for C++ compilation.This updates the XLA integration to match current JAX FFI expectations while keeping the header dependency sourced from upstream GitHub code.
Technical Details
xla/ffi/apisubtree from the XLA source commit pinned byjax-v0.9.2.envpool/core/xla_template.h: rewrites the custom call bridge to usexla::ffityped handlers for CPU and GPU.envpool/python/xla_template.py: registers API version 1 FFI targets and threads the custom handle through attrs with input/output aliasing.envpool/workspace0.bzl: adds the pinnedopenxla_ffi_headerssource archive.third_party/openxla_ffi/ffi_api.BUILD: exposes only the exported FFI headers with the expected include prefix.openxla/xla@187a5eb58277a85847d1516bd1e20b7faf03d5ef, which is the XLA revision referenced byjax-v0.9.2.Test Plan
Automated
brix ssh dev-0 -C -- 'cd /root/code/envpool && USE_BAZEL_VERSION=8.6.0 bazel --output_user_root=/tmp/envpool-openxla-source-subtree-bazel build //envpool/core:py_envpool --@rules_python//python/config_settings:python_version=3.11': passedbrix ssh dev-0 -C -- 'cd /root/code/envpool && USE_BAZEL_VERSION=8.6.0 bazel --output_user_root=/tmp/envpool-openxla-source-subtree-bazel build //envpool/core:py_envpool --@rules_python//python/config_settings:python_version=3.12': passedbrix ssh dev-0 -C -- 'cd /root/code/envpool && USE_BAZEL_VERSION=8.6.0 bazel --output_user_root=/tmp/envpool-openxla-source-subtree-bazel build //envpool/core:py_envpool --@rules_python//python/config_settings:python_version=3.13': passedbrix ssh dev-0 -C -- 'cd /root/code/envpool && BAZELOPT="--output_user_root=/tmp/envpool-openxla-subtree-wheel-bazel2" make bazel-build': passedbrix ssh dev-0 -C -- 'cd /root/code/envpool && /tmp/envpool-issue303-fix-venv/bin/python examples/xla_step.py': passed earlier during the typed FFI migration validation ondev-0Suggested Manual
pip install <built wheel> "jax[cuda12]" gym packaging: verify install against the same JAX family used in the repro.python examples/xla_step.py: confirm the GPU path runs without the oldAPI_VERSION_ORIGINALwarning.