XLA Custom Calls

This document describes how to write and use XLA custom calls using XLA FFI library. Custom call is a mechanism to describe an external "operation" in the HLO module to the XLA compiler (at compile time), and XLA FFI is a mechanism to register implementation of such operations with XLA (at run time). FFI stands for "foreign function interface" and it is a set of C APIs that define a binary interface (ABI) for XLA to call into external code written in other programming languages. XLA provides header-only bindings for XLA FFI written in C++, which hides all the low level details of underlying C APIs from the end user.

JAX + XLA Custom Calls

See JAX documentation for end to end examples of integrating custom calls and XLA FFI with JAX.

XLA FFI Binding

XLA FFI binding is a compile-time specification of the custom call signature: custom call arguments, attributes and their types, and additional parameters passed via the execution context (i.e., gpu stream for GPU backend). XLA FFI binding can be bound to any C++ callable (function pointer, lambda, etc.) with compatible operator() signature. Constructed handler decodes XLA FFI call frame (defined by the stable C API), type check all parameters, and forward decoded results to the user-defined callback.

XLA FFI binding heavily relies on template metaprogramming to be be able to compile constructed handler to the most efficient machine code. Run time overheads are in order of a couple of nanoseconds for each custom call parameter.

XLA FFI customization points implemented as template specializations, and users can define how to decode their custom types, i.e., it is possible to define custom decoding for user-defined enum class types.

Returning Errors From Custom Calls

Custom call implementations must return xla::ffi::Error value to signal success or error to XLA runtime. It is similar to absl::Status, and has the same set of error codes. We do not use absl::Status because it does not have a stable ABI and it would be unsafe to pass it between dynamically loaded custom call library, and XLA itself.

// Handler that always returns an error.
auto always_error = Ffi::Bind().To(
    []() { return Error(ErrorCode::kInternal, "Oops!"); });

// Handler that always returns a success.
auto always_success = Ffi::Bind().To(
    []() { return Error::Success(); });

Buffer Arguments And Results

XLA uses destination passing style for results: custom calls (or any other XLA operations for that matter) do not allocate memory for results, and instead write into destinations passed by XLA runtime. XLA uses static buffer assignment, and allocates buffers for all values based on their live ranges at compile time.

Results passed to FFI handlers wrapped into a Result<T> template, that has a pointer-like semantics: operator-> gives access to the underlying parameter.

AnyBuffer arguments and results gives access to custom call buffer parameters of any data type. This is useful when custom call has a generic implementation that works for multiple data types, and custom call implementation does run time dispatching based on data type. AnyBuffer gives access to the buffer data type, dimensions, and a pointer to the buffer itself.

%0 = "stablehlo.custom_call"(%arg0) {
  call_target_name = "foo",
  api_version = 4 : i32
} : (tensor<2x2xf32>) -> tensor<2x2xf32>

// Buffers of any number of dimensions and data type.
auto handler = Ffi::Bind().Arg<AnyBuffer>().Ret<AnyBuffer>().To(
    [](AnyBuffer arg, Result<AnyBuffer> res) -> Error {
      void* arg_data = arg.untyped_data();
      void* res_data = res->untyped_data();
      return Error::Success();
    });

Constrained Buffer Arguments And Results

Buffer allows to add constraints on the buffer data type and number of dimensions, and they will be automatically checked by the handler and return an error to XLA runtime, if run time arguments do not match the FFI handler signature.

// Buffers of any number of dimensions and F32 data type.
auto handler = Ffi::Bind().Arg<Buffer<F32>>().Ret<Buffer<F32>>().To(
    [](Buffer<F32> arg, Result<Buffer<F32>> res) -> Error {
      float* arg_data = arg.typed_data();
      float* res_data = res->typed_data();
      return Error::Success();
    });

// Buffers of number of dimensions 2 and F32 data type.
auto handler = Ffi::Bind().Arg<BufferR2<F32>>().Ret<BufferR2<F32>>().To(
    [](BufferR2<F32> arg, Result<BufferR2<F32>> res) -> Error {
      float* arg_data = arg.typed_data();
      float* res_data = res->typed_data();
      return Error::Success();
    });

Matching And Verifying Buffers

Buffer patterns can express constraints that are more specific than the argument and result types in an FFI binding. They are useful for refining an AnyBuffer to a concrete buffer type, checking dimension sizes, and checking relationships between multiple buffer shapes.

Patterns are immutable and can specify dtype and rank either directly or with the corresponding modifiers:

namespace m =