Enable XLA on MacOS and fix XLA on GPUs (#707)

Author: pseudo-rnd-thoughts

.pre-commit-config.yaml

@@ -36,17 +36,17 @@ repos:
           - --show-source
           - --statistics
   - repo: https://github.com/asottile/pyupgrade
-    rev: v3.21.1
+    rev: v3.21.2
     hooks:
       - id: pyupgrade
         args: ["--py310-plus"]
   - repo: https://github.com/PyCQA/isort
-    rev: 7.0.0
+    rev: 9.0.0a3
     hooks:
       - id: isort
         args: ["--profile", "black"]
   - repo: https://github.com/python/black
-    rev: 25.9.0
+    rev: 26.5.1
     hooks:
       - id: black
   - repo: https://github.com/pycqa/pydocstyle

pyproject.toml

@@ -2,7 +2,7 @@
 requires = [
     "scikit-build-core>=0.10",
     "nanobind>=2.5.0",
-    "jax >= 0.4.31; sys_platform == 'win32' or sys_platform == 'linux'"
+    "jax >= 0.4.31"
 ]
 build-backend = "scikit_build_core.build"
 
@@ -50,15 +50,15 @@ vector = [
 xla = [
     "gymnasium>=1.1.0",
     "opencv-python>=3.0",
-    "jax >= 0.4.31; sys_platform == 'win32' or sys_platform == 'linux'",
-    "chex; sys_platform == 'win32' or sys_platform == 'linux'",
+    "jax >= 0.4.31",
+    "chex",
 ]
 test = [
     "pytest>=7.0",
     "gymnasium>=1.1.0",
     "opencv-python>=3.0",
-    "jax >= 0.4.31; sys_platform == 'win32' or sys_platform == 'linux'",
-    "chex; sys_platform == 'win32' or sys_platform == 'linux'"
+    "jax >= 0.4.31",
+    "chex"
 ]
 
 [project.urls]
@@ -73,29 +73,26 @@ testpaths = ["tests"]
 norecursedirs = ["vcpkg", "build"]
 
 [tool.scikit-build]
+# XLA is enabled by default which requires jax to be installed
 cmake.args = [
     "-DSDL_SUPPORT=ON",
     "-DSDL_DYNLOAD=ON",
     "-DBUILD_CPP_LIB=OFF",
     "-DBUILD_PYTHON_LIB=ON",
-    "-DBUILD_VECTOR_LIB=ON"
+    "-DBUILD_VECTOR_LIB=ON",
+    "-DBUILD_VECTOR_XLA_LIB=ON"
 ]
 cmake.build-type = "Release"
 metadata.version.provider = "scikit_build_core.metadata.regex"
 metadata.version.input = "version.txt"
 metadata.version.regex = "(?P<value>.+)"
 
-# Enable XLA on Windows and Linux
+# On windows, vcpkg's default triplet (x64-windows) links dynamically but SDL2
+# needs to be static, therefore, force the custom triplet. Linux/macOS
+# are unaffected because vcpkg defaults to static linkage there.
 [[tool.scikit-build.overrides]]
-if.platform-system = "^(win32|linux)"
-cmake.args = [
-    "-DSDL_SUPPORT=ON",
-    "-DSDL_DYNLOAD=ON",
-    "-DBUILD_CPP_LIB=OFF",
-    "-DBUILD_PYTHON_LIB=ON",
-    "-DBUILD_VECTOR_LIB=ON",
-    "-DBUILD_VECTOR_XLA_LIB=ON"  # add vector-xla-lib="ON"
-]
+if.platform-system = "^win32"
+cmake.define.VCPKG_TARGET_TRIPLET = "x64-windows-mixed"
 
 [tool.cibuildwheel]
 # Skip 32-bit wheels, PyPy, musllinux, cp313t (free-threading was experimental in 3.13),

scripts/benchmark-vec-env.py

@@ -3,6 +3,7 @@
 
 Allows testing of the impact of `num_envs`, `batch_size`, `autoreset_mode`, and `num_threads` on steps per second and resource usage.
 """
+
 from __future__ import annotations
 
 import dataclasses

src/ale/python/CMakeLists.txt

@@ -22,24 +22,37 @@ if (BUILD_VECTOR_LIB AND BUILD_VECTOR_XLA_LIB)
     )
     message(STATUS "XLA include directory: '${XLA_DIR}'")
 
-    # Find CUDA toolkit for headers and linking (optional - XLA disabled if not found)
+    # CUDA toolkit is optional: when found we additionally compile/link the GPU
+    # handlers, otherwise XLA is built with CPU-only support (e.g. macOS).
     find_package(CUDAToolkit QUIET)
 
+    # Always build the CPU XLA handlers. The GPU handlers in
+    # ale_vector_xla_interface.cpp are guarded by BUILD_VECTOR_XLA_GPU, which is
+    # only defined below when the CUDA Toolkit is found, so they are naturally
+    # skipped when CUDA is unavailable (e.g. macOS).
+    target_sources(_ale_py PRIVATE ale_vector_python_interface.cpp ale_vector_xla_interface.cpp)
+    target_include_directories(_ale_py PUBLIC ${XLA_DIR})
+    target_compile_definitions(_ale_py PRIVATE BUILD_VECTOR_LIB BUILD_VECTOR_XLA_LIB)
+    set_target_properties(_ale_py PROPERTIES
+        POSITION_INDEPENDENT_CODE ON
+    )
+
     if (CUDAToolkit_FOUND)
-        message(STATUS "CUDA Toolkit found: ${CUDAToolkit_VERSION}")
-        target_sources(_ale_py PRIVATE ale_vector_python_interface.cpp ale_vector_xla_interface.cpp)
-        target_include_directories(_ale_py PUBLIC ${XLA_DIR})
+        message(STATUS "CUDA Toolkit found: ${CUDAToolkit_VERSION} - enabling GPU XLA support")
+        # The GPU handlers only use the CUDA runtime API (no device kernels), so
+        # the host compiler can build them given the CUDA headers and cudart.
         target_include_directories(_ale_py PRIVATE ${CUDAToolkit_INCLUDE_DIRS})
-        target_compile_definitions(_ale_py PRIVATE BUILD_VECTOR_LIB BUILD_VECTOR_XLA_LIB)
-        target_link_libraries(_ale_py PRIVATE CUDA::cudart)
-
-        set_target_properties(_ale_py PROPERTIES
-            POSITION_INDEPENDENT_CODE ON
-        )
+        target_compile_definitions(_ale_py PRIVATE BUILD_VECTOR_XLA_GPU)
+        if (WIN32)
+            # Static linking bakes cudart into _ale_py.pyd so there is no runtime DLL to find.
+            target_link_libraries(_ale_py PRIVATE CUDA::cudart_static)
+        else()
+            # The loader resolves libcudart via the toolkit's lib path / the manylinux image, and
+            # jax already loads the same runtime in-process, so there is no missing-library issue.
+            target_link_libraries(_ale_py PRIVATE CUDA::cudart)
+        endif()
     else()
-        message(WARNING "CUDA Toolkit not found - disabling XLA support. Install CUDA Toolkit to enable XLA.")
-        target_sources(_ale_py PRIVATE ale_vector_python_interface.cpp)
-        target_compile_definitions(_ale_py PRIVATE BUILD_VECTOR_LIB)
+        message(STATUS "CUDA Toolkit not found - building CPU-only XLA support")
     endif()
 elseif (BUILD_VECTOR_LIB)
     target_sources(_ale_py PRIVATE ale_vector_python_interface.cpp)

src/ale/python/__init__.py

@@ -17,10 +17,8 @@
         ctypes.CDLL("vcruntime140.dll")
         ctypes.CDLL("msvcp140.dll")
     except OSError:
-        raise OSError(
-            """Microsoft Visual C++ Redistribution Pack is not installed.
-It can be downloaded from https://aka.ms/vs/16/release/vc_redist.x64.exe."""
-        )
+        raise OSError("""Microsoft Visual C++ Redistribution Pack is not installed.
+It can be downloaded from https://aka.ms/vs/16/release/vc_redist.x64.exe.""")
 
     # Loading DLLs on Windows is kind of a disaster
     # The best approach seems to be using LoadLibraryEx with user defined search paths.

src/ale/python/ale_vector_xla_interface.cpp

@@ -9,7 +9,7 @@
 #include <cstring>  // For memcpy
 #include <iostream>
 
-#ifdef __CUDACC__
+#ifdef BUILD_VECTOR_XLA_GPU
 #include <cuda_runtime.h>
 #endif
 
@@ -108,7 +108,7 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(
         .Ret<ffi::Buffer<ffi::S32>>()  // episode_frame_numbers
 );
 
-#ifdef __CUDACC__
+#ifdef BUILD_VECTOR_XLA_GPU
 // GPU version of XLAReset with CUDA stream support
 ffi::Error XLAResetGPUImpl(
     cudaStream_t stream,
@@ -270,7 +270,7 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(
         .Ret<ffi::Buffer<ffi::S32>>()  // frame_numbers
         .Ret<ffi::Buffer<ffi::S32>>()  // episode_frame_numbers
 );
-#endif  // __CUDACC__
+#endif  // BUILD_VECTOR_XLA_GPU
 
 // CPU version of XLAStep
 ffi::Error XLAStepImpl(
@@ -394,7 +394,7 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(
         .Ret<ffi::Buffer<ffi::S32>>()   // episode_frame_numbers
 );
 
-#ifdef __CUDACC__
+#ifdef BUILD_VECTOR_XLA_GPU
 // GPU version of XLAStep with CUDA stream support
 ffi::Error XLAStepGPUImpl(
     cudaStream_t stream,
@@ -610,7 +610,7 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(
         .Ret<ffi::Buffer<ffi::S32>>()   // frame_numbers
         .Ret<ffi::Buffer<ffi::S32>>()   // episode_frame_numbers
 );
-#endif  // __CUDACC__
+#endif  // BUILD_VECTOR_XLA_GPU
 
 template <typename T>
 nb::capsule EncapsulateFFICall(T *fn) {
@@ -625,7 +625,7 @@ void init_vector_module_xla(nb::module_& m) {
     m.def("VectorXLAReset", [] {return EncapsulateFFICall(AtariVectorEnvXLAReset); });
     m.def("VectorXLAStep", [] {return EncapsulateFFICall(AtariVectorEnvXLAStep); });
 
-#ifdef __CUDACC__
+#ifdef BUILD_VECTOR_XLA_GPU
     // GPU handlers
     m.def("VectorXLAResetGPU", [] {return EncapsulateFFICall(AtariVectorEnvXLAResetGPU); });
     m.def("VectorXLAStepGPU", [] {return EncapsulateFFICall(AtariVectorEnvXLAStepGPU); });

src/ale/python/env.py

@@ -189,7 +189,7 @@ def __init__(
                 low=0, high=255, dtype=np.uint8, shape=(self.ale.getRAMSize(),)
             )
         elif self._obs_type == "rgb" or self._obs_type == "grayscale":
-            (screen_height, screen_width) = self.ale.getScreenDims()
+            screen_height, screen_width = self.ale.getScreenDims()
             image_shape = (
                 screen_height,
                 screen_width,

tests/python/test_atari_vector_xla.py

@@ -130,3 +130,55 @@ def actor(handle, action):
 
     envs_1.close()
     envs_2.close()
+
+
+CPU_DEVICES = jax.devices("cpu")
+try:
+    GPU_DEVICES = jax.devices("gpu")
+except RuntimeError:
+    GPU_DEVICES = []
+
+
+def is_gpu_array(arr):
+    return all(arr_device in GPU_DEVICES for arr_device in arr.devices())
+
+
+def is_cpu_array(arr):
+    return all(arr_device in CPU_DEVICES for arr_device in arr.devices())
+
+
+def test_xla_dispatch_gpu():
+    envs = AtariVectorEnv("pong", num_envs=1)
+    handle, xla_reset, xla_step = envs.xla()
+
+    print(f"GPU_DEVICES={GPU_DEVICES}, CPU_DEVICES={CPU_DEVICES}")
+    if GPU_DEVICES:
+        print("RUNNING ON GPUs")
+        handle, (obs, info) = xla_reset(handle)
+        assert is_gpu_array(obs)
+        assert is_gpu_array(info["frame_number"])
+
+        actions = np.zeros(1, dtype=np.int32)
+        handle, (obs, rewards, terminations, truncations, info) = xla_step(
+            handle, actions
+        )
+        assert is_gpu_array(obs)
+        assert is_gpu_array(rewards)
+        assert is_gpu_array(terminations)
+        assert is_gpu_array(truncations)
+        assert is_gpu_array(info["frame_number"])
+    else:
+        print("RUNNING ON CPUs")
+        handle, (obs, info) = xla_reset(handle)
+        assert is_cpu_array(obs)
+        assert is_cpu_array(info["frame_number"])
+
+        actions = np.zeros(1, dtype=np.int32)
+        handle, (obs, rewards, terminations, truncations, info) = xla_step(
+            handle, actions
+        )
+        assert is_cpu_array(obs)
+        assert is_cpu_array(rewards)
+        assert is_cpu_array(terminations)
+        assert is_cpu_array(truncations)
+        assert is_cpu_array(info["frame_number"])