[envpool] migrate XLA custom calls to OpenXLA FFI

envpool/core/BUILD

@@ -196,6 +196,7 @@ cc_library(
     ],
     deps = [
         ":spec",
+        "@openxla_ffi_headers//:headers",
         "@cuda//:cudart_static",
         "@pybind11",
     ],

envpool/core/xla_template.h

@@ -22,6 +22,7 @@
 #include <pybind11/numpy.h>
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
+#include <xla/ffi/api/ffi.h>
 
 #include <array>
 #include <cstddef>
@@ -35,25 +36,13 @@
 #include "envpool/core/spec.h"
 
 namespace py = pybind11;
+namespace xla_ffi = xla::ffi;
 
 template <typename Spec>
 static auto SpecToTuple(const Spec& spec) {
   return std::make_tuple(py::dtype::of<typename Spec::dtype>(), spec.shape);
 }
 
-template <std::size_t N>
-void ToArray(const void** raw, std::array<void*, N>* array) {
-  int i = 0;
-  std::apply([&](auto&&... a) { ((a = const_cast<void*>(raw[i++])), ...); },
-             *array);
-}
-
-template <std::size_t N>
-void ToArray(void** raw, std::array<void*, N>* array) {
-  int i = 0;
-  std::apply([&](auto&&... a) { ((a = raw[i++]), ...); }, *array);
-}
-
 template <typename Class, typename CC>
 struct CustomCall {
   using InSpecs = std::invoke_result_t<decltype(CC::InSpecs), Class*>;
@@ -66,35 +55,98 @@ struct CustomCall {
         std::string(reinterpret_cast<const char*>(&obj), sizeof(Class*)));
   }
 
-  static void Cpu(void* out, const void** in) {
-    Class* obj = nullptr;
-    std::memcpy(reinterpret_cast<void*>(&obj), in[0], sizeof(Class*));
-    in += 1;
-    In in_arr;
-    Out out_arr;
-    ToArray(in, &in_arr);
-    if (std::tuple_size<Out>::value == 0) {
-      std::memcpy(out, reinterpret_cast<const void*>(&obj), sizeof(Class*));
-    } else {
-      void** outs = reinterpret_cast<void**>(out);
-      std::memcpy(outs[0], reinterpret_cast<const void*>(&obj), sizeof(Class*));
-      ToArray(outs + 1, &out_arr);
+  static xla_ffi::ErrorOr<Class*> ResolveHandle(xla_ffi::Dictionary attrs) {
+    auto handle = attrs.get<std::int64_t>("handle");
+    if (!handle) {
+      return xla_ffi::Unexpected(handle.error());
+    }
+    return reinterpret_cast<Class*>(
+        static_cast<std::uintptr_t>(static_cast<std::int64_t>(*handle)));
+  }
+
+  static xla_ffi::Error ValidateArity(xla_ffi::RemainingArgs args,
+                                      xla_ffi::RemainingRets rets) {
+    constexpr std::size_t kExpectedArgs = std::tuple_size_v<In> + 1;
+    constexpr std::size_t kExpectedRets = std::tuple_size_v<Out> + 1;
+    if (args.size() != kExpectedArgs) {
+      return xla_ffi::Error::InvalidArgument(
+          "Expected " + std::to_string(kExpectedArgs) + " buffers, got " +
+          std::to_string(args.size()));
+    }
+    if (rets.size() != kExpectedRets) {
+      return xla_ffi::Error::InvalidArgument(
+          "Expected " + std::to_string(kExpectedRets) + " results, got " +
+          std::to_string(rets.size()));
+    }
+    return xla_ffi::Error();
+  }
+
+  static xla_ffi::Error PopulateInBuffers(xla_ffi::RemainingArgs args,
+                                          In* in_arr) {
+    for (std::size_t i = 0; i < in_arr->size(); ++i) {
+      auto buffer = args.get<xla_ffi::AnyBuffer>(i + 1);
+      if (!buffer) {
+        return buffer.error();
+      }
+      (*in_arr)[i] = (*buffer).untyped_data();
+    }
+    return xla_ffi::Error();
+  }
+
+  static xla_ffi::Error PopulateOutBuffers(xla_ffi::RemainingRets rets,
+                                           Out* out_arr) {
+    for (std::size_t i = 0; i < out_arr->size(); ++i) {
+      auto buffer = rets.get<xla_ffi::AnyBuffer>(i + 1);
+      if (!buffer) {
+        return buffer.error();
+      }
+      (*out_arr)[i] = (*buffer)->untyped_data();
+    }
+    return xla_ffi::Error();
+  }
+
+  static xla_ffi::Error CpuExecute(xla_ffi::RemainingArgs args,
+                                   xla_ffi::RemainingRets rets,
+                                   xla_ffi::Dictionary attrs) {
+    if (auto err = ValidateArity(args, rets); err.failure()) {
+      return err;
+    }
+    auto obj = ResolveHandle(attrs);
+    if (!obj) {
+      return obj.error();
     }
-    CC::Cpu(obj, in_arr, out_arr);
+    In in_arr{};
+    Out out_arr{};
+    if (auto err = PopulateInBuffers(args, &in_arr); err.failure()) {
+      return err;
+    }
+    if (auto err = PopulateOutBuffers(rets, &out_arr); err.failure()) {
+      return err;
+    }
+    CC::Cpu(*obj, in_arr, out_arr);
+    return xla_ffi::Error();
   }
 
-  static void Gpu(cudaStream_t stream, void** buffers, const char* opaque,
-                  std::size_t opaque_len) {
-    Class* obj = nullptr;
-    std::memcpy(reinterpret_cast<void*>(&obj), opaque, sizeof(Class*));
-    buffers += 1;
-    In in_arr;
-    Out out_arr;
-    ToArray(buffers, &in_arr);
-    buffers += std::tuple_size<In>::value;
-    buffers += 1;
-    ToArray(buffers, &out_arr);
-    CC::Gpu(obj, stream, in_arr, out_arr);
+  static xla_ffi::Error GpuExecute(
+      cudaStream_t stream, xla_ffi::RemainingArgs args,
+      xla_ffi::RemainingRets rets, xla_ffi::Dictionary attrs) {
+    if (auto err = ValidateArity(args, rets); err.failure()) {
+      return err;
+    }
+    auto obj = ResolveHandle(attrs);
+    if (!obj) {
+      return obj.error();
+    }
+    In in_arr{};
+    Out out_arr{};
+    if (auto err = PopulateInBuffers(args, &in_arr); err.failure()) {
+      return err;
+    }
+    if (auto err = PopulateOutBuffers(rets, &out_arr); err.failure()) {
+      return err;
+    }
+    CC::Gpu(*obj, stream, in_arr, out_arr);
+    return xla_ffi::Error();
   }
 
   static auto Specs(Class* obj) {
@@ -113,9 +165,22 @@ struct CustomCall {
   }
 
   static auto Capsules() {
+    XLA_FFI_DEFINE_HANDLER(
+        cpu_handler, CpuExecute,
+        xla_ffi::Ffi::Bind()
+            .RemainingArgs()
+            .RemainingRets()
+            .Attrs<xla_ffi::Dictionary>());
+    XLA_FFI_DEFINE_HANDLER(
+        gpu_handler, GpuExecute,
+        xla_ffi::Ffi::Bind()
+            .Ctx<xla_ffi::PlatformStream<cudaStream_t>>()
+            .RemainingArgs()
+            .RemainingRets()
+            .Attrs<xla_ffi::Dictionary>());
     return std::make_tuple(
-        py::capsule(reinterpret_cast<void*>(Cpu), "xla._CUSTOM_CALL_TARGET"),
-        py::capsule(reinterpret_cast<void*>(Gpu), "xla._CUSTOM_CALL_TARGET"));
+        py::capsule(reinterpret_cast<void*>(cpu_handler)),
+        py::capsule(reinterpret_cast<void*>(gpu_handler)));
   }
 
   static auto Xla(Class* obj) {

envpool/python/xla_template.py

@@ -14,7 +14,8 @@
 """xla template on python side."""
 
 from collections import namedtuple
-from typing import Any, Callable, cast
+import sys
+from typing import Any, Callable
 
 import numpy as np
 from jax import ShapeDtypeStruct, dtypes, ffi
@@ -53,13 +54,13 @@ def _make_xla_function(
         call_target_name,
         cpu_capsule,
         platform="cpu",
-        api_version=0,
+        api_version=1,
     )
     ffi.register_ffi_target(
         call_target_name,
         gpu_capsule,
         platform="gpu",
-        api_version=0,
+        api_version=1,
     )
     result_specs = tuple(_shape_dtype_struct(*spec) for spec in out_specs)
     xla_func = ffi.ffi_call(
@@ -72,14 +73,12 @@ def _make_xla_function(
             if len(out_specs) > 1
             else _layout(out_specs[0][0])
         ),
-        # JAX target registration uses api_version=0 for the legacy untyped
-        # handler, but StableHLO custom_call uses API_VERSION_ORIGINAL == 1.
-        custom_call_api_version=1,
-        legacy_backend_config=cast(Any, handle),
+        input_output_aliases={0: 0},
     )
+    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)
 
     return call
 

envpool/workspace0.bzl

@@ -117,6 +117,17 @@ def workspace():
         ],
     )
 
+    maybe(
+        http_archive,
+        name = "openxla_ffi_headers",
+        build_file = "//third_party/openxla_ffi:ffi_api.BUILD",
+        sha256 = "753df38eab0d430da20e614316401663bcfca433905b976745a6e59998635ce8",
+        strip_prefix = "xla-187a5eb58277a85847d1516bd1e20b7faf03d5ef/xla/ffi/api",
+        urls = [
+            "https://github.com/openxla/xla/archive/187a5eb58277a85847d1516bd1e20b7faf03d5ef.tar.gz",
+        ],
+    )
+
     maybe(
         http_archive,
         name = "com_google_absl",

third_party/openxla_ffi/BUILD

@@ -0,0 +1,17 @@
+# Copyright 2022 Garena Online Private Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+package(default_visibility = ["//visibility:public"])
+
+exports_files(["ffi_api.BUILD"])

third_party/openxla_ffi/README.md

@@ -0,0 +1,12 @@
+# OpenXLA FFI headers
+
+Pinned OpenXLA FFI headers used by `envpool/core/xla_template.h`.
+
+Source archive:
+
+- `openxla/xla@187a5eb58277a85847d1516bd1e20b7faf03d5ef`
+- fetched via `http_archive` in `envpool/workspace0.bzl`
+- this is the XLA revision pinned by `jax-v0.9.2` in `third_party/xla/revision.bzl`
+- only the `xla/ffi/api/` subtree is extracted into the external repo
+
+Only `xla/ffi/api/{api.h,c_api.h,ffi.h}` is exposed to Bazel.

third_party/openxla_ffi/ffi_api.BUILD

@@ -0,0 +1,27 @@
+# Copyright 2022 Garena Online Private Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+load("@rules_cc//cc:defs.bzl", "cc_library")
+
+package(default_visibility = ["//visibility:public"])
+
+cc_library(
+    name = "headers",
+    hdrs = [
+        "api.h",
+        "c_api.h",
+        "ffi.h",
+    ],
+    include_prefix = "xla/ffi/api",
+)