Llamadas personalizadas de XLA

En este documento, se describe cómo escribir y usar llamadas personalizadas de XLA con la biblioteca XLA FFI. Una llamada personalizada es un mecanismo para describir una "operación" externa en el módulo HLO al compilador de XLA (en el tiempo de compilación), y XLA FFI es un mecanismo para registrar la implementación de esas operaciones con XLA (en el tiempo de ejecución). FFI significa "interfaz de función externa" y es un conjunto de APIs de C que definen una interfaz binaria (ABI) para que XLA llame a código externo escrito en otros lenguajes de programación. XLA proporciona vinculaciones solo de encabezado para XLA FFI escritas en C++, que oculta todos los detalles de bajo nivel de las APIs de C subyacentes del usuario final.

Llamadas personalizadas de JAX + XLA

Consulta la documentación de JAX para obtener ejemplos de extremo a extremo de la integración de llamadas personalizadas y XLA FFI con JAX.

Vinculación de XLA FFI

La vinculación de XLA FFI es una especificación en el momento de la compilación de la firma de llamada personalizada: argumentos de llamada personalizada, atributos y sus tipos, y parámetros adicionales que se pasan a través del contexto de ejecución (es decir, transmisión de GPU para el backend de GPU). La vinculación de XLA FFI se puede vincular a cualquier elemento invocable de C++ (puntero de función, lambda, etcétera) con una firma operator() compatible. El controlador construido decodifica el marco de llamada de XLA FFI (definido por la API de C estable), verifica el tipo de todos los parámetros y reenvía los resultados decodificados a la devolución de llamada definida por el usuario.

La vinculación de XLA FFI se basa en gran medida en la metaprogramación de plantillas para poder compilar el controlador construido en el código máquina más eficiente. Los gastos generales del tiempo de ejecución son del orden de un par de nanosegundos para cada parámetro de llamada personalizada.

Los puntos de personalización de XLA FFI se implementan como especializaciones de plantillas, y los usuarios pueden definir cómo decodificar sus tipos personalizados. Es decir, es posible definir la decodificación personalizada para los tipos enum class definidos por el usuario.

Cómo mostrar errores de llamadas personalizadas

Las implementaciones de llamadas personalizadas deben mostrar el valor xla::ffi::Error para indicar éxito o error en el tiempo de ejecución de XLA. Es similar a absl::Status y tiene el mismo conjunto de códigos de error. No usamos absl::Status porque no tiene una ABI estable y no sería seguro pasarla entre la biblioteca de llamadas personalizadas cargada de forma dinámica y XLA.

// 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(); });

Argumentos y resultados del búfer

XLA usa el estilo de paso de destino para los resultados: las llamadas personalizadas (o cualquier otra operación de XLA) no asignan memoria para los resultados y, en su lugar, escriben en los destinos que pasa el tiempo de ejecución de XLA. XLA usa la asignación de búfer estático y asigna búferes para todos los valores según sus rangos activos en el tiempo de compilación.

Los resultados que se pasan a los controladores de FFI se encapsulan en una plantilla Result<T>, que tiene una semántica similar a la de un puntero: operator-> da acceso al parámetro subyacente.

Los argumentos y resultados de AnyBuffer dan acceso a los parámetros de búfer de llamada personalizada de cualquier tipo de datos. Esto es útil cuando la llamada personalizada tiene una implementación genérica que funciona para varios tipos de datos y la implementación de llamada personalizada ejecuta el envío en el tiempo de ejecución según el tipo de datos. AnyBuffer da acceso al tipo de datos del búfer, las dimensiones y un puntero al búfer.

%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();
    });

Argumentos y resultados del búfer restringido

Buffer permite agregar restricciones en el tipo de datos del búfer y la cantidad de dimensiones, y el controlador las verificará automáticamente y mostrará un error en el tiempo de ejecución de XLA si los argumentos del tiempo de ejecución no coinciden con la firma del controlador de FFI.

// 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();
    });

Cómo hacer coincidir y verificar búferes

Los patrones de búfer pueden expresar restricciones que son más específicas que los tipos de argumentos y resultados en una vinculación de FFI. Son útiles para refinar un AnyBuffer a un tipo de búfer concreto, verificar los tamaños de las dimensiones y verificar las relaciones entre varias formas de búfer.

Los patrones son inmutables y pueden especificar dtype y rango directamente o con los modificadores correspondientes:

namespace m = ::xla::ffi::match;

m::Buffer<F32, 2>();
m::Buffer().WithDType<F32>().WithRank<2>();

La función Match verifica un patrón y muestra su tipo de búfer con la mayor restricción. Para refinar un AnyBuffer, el patrón debe especificar exactamente un tipo de datos y un rango. Hacer coincidir un búfer ya escrito verifica cualquier restricción adicional y conserva su tipo.

namespace m = ::xla::ffi::match;

auto handler = Ffi::Bind().Arg<AnyBuffer>().Ret<AnyBuffer>().To(
    [](AnyBuffer input, Result<AnyBuffer> output) -> Error {
      int64_t rows;
      int64_t cols;

      ErrorOr<BufferR2<F32>> typed_input =
          Match("input", input,
                m::Buffer<F32>().WithDims(&rows, &cols));
      if (typed_input.has_error()) {
        return std::move(typed_input).error();
      }

      ErrorOr<Result<BufferR2<F32>>> typed_output =
          Match("output", output,
                m::Buffer<F32>().WithDims(rows, cols));
      if (typed_output.has_error()) {
        return std::move(typed_output).error();
      }

      float* input_data = typed_input->typed_data();
      float* output_data = (*typed_output)->typed_data();
      return Error::Success();
    });

En este ejemplo, la coincidencia input captura sus dos tamaños de dimensión, y la coincidencia output verifica que tenga la misma forma. Las capturas se confirman solo cuando el patrón completo coincide correctamente.

Un patrón de dimensión puede no estar restringido, ser fijo o capturado. Los valores fijos y los punteros de captura se convierten de forma implícita en patrones de dimensión:

m::Dim()       // Any dimension size.
m::Dim(16)     // A dimension whose size is 16.
m::Dim(&rows)  // Any dimension size, captured in `rows` on success.

m::Buffer<F32>().WithDims(&rows, 16);

WithDims describe una forma completa y establece su rango a partir de la cantidad de argumentos. WithDim<I> o WithDim(index, ...) restringe las dimensiones de forma posicional sin fijar el rango completo. Una restricción posicional requiere que exista la dimensión a la que se hace referencia, pero deja sin restricciones todas las demás dimensiones:

ErrorOr<BufferR4<F32>> MatchActivations(AnyBuffer input) {
  return Match(
      "input", input,
      m::Buffer<F32, 4>().WithDim<0>(1).WithDim<3>(128));
}

Error VerifyDimension(AnyBuffer input, size_t index, int64_t size) {
  return Verify("input", input, m::Buffer().WithDim(index, size));
}

MatchActivations requiere un búfer F32 de rango 4 con dimensiones 0 y 3 iguales a 1 y 128, respectivamente; las dimensiones 1 y 2 no están restringidas. Sin un explícito WithRank, WithDim(index, ...) acepta cualquier rango mayor que index.

Las capturas de dimensión también pueden expresar relaciones entre posiciones:

auto matrix = m::Buffer<F32>().WithDims(&rows, 128);

int64_t n;
ErrorOr<BufferR2<F32>> square =
    Match("matrix", buffer,
          m::Buffer<F32>().WithDims(&n, &n));

Cuando el mismo puntero de captura aparece varias veces, todas las dimensiones correspondientes deben tener el mismo tamaño. Por lo tanto, la llamada Match anterior solo acepta matrices cuadradas. Las capturas de dimensión se escriben solo después de que el patrón completo se realiza correctamente y no se modifican en caso de falla.

Las relaciones de búfer completo no requieren la captura de cada dimensión. WithShapeOf coincide con la forma completa del tiempo de ejecución de otro búfer y permite un dtype diferente. Like también requiere el mismo dtype:

Error VerifySortBuffers(AnyBuffer keys, AnyBuffer values,
                        Result<AnyBuffer> keys_output) {
  Error error =
      Verify("values", values, m::Buffer().WithShapeOf(keys));
  if (error.failure()) {
    return error;
  }

  return Verify("keys output", *keys_output, m::Buffer().Like(keys));
}

Ambos modificadores copian los metadatos del búfer de referencia cuando se compila el patrón; el patrón no conserva una referencia al búfer. Debido a que estas son restricciones del tiempo de ejecución, no refinan un AnyBuffer a un tipo de datos que se devuelve concreto. Son útiles principalmente con Verify o con Match cuando el búfer de entrada ya tiene un tipo concreto.

Usa Verify cuando el búfer deba conservar su tipo existente o cuando un patrón acepte más de un tipo de datos o rango. Por ejemplo, un búfer de índice puede aceptar S32 o S64 y enviar en su tipo de datos después de la verificación:

namespace m = ::xla::ffi::match;

Error VerifyIndices(AnyBuffer indices) {
  auto pattern = m::Buffer().WithDType<S32, S64>().WithRank<1, 2>();

  Error error = Verify("indices", indices, pattern);
  if (error.failure()) {
    return error;
  }

  // Dispatch an implementation based on `indices.element_type()`.
  return Error::Success();
}

Debido a que este patrón permite varios tipos de búfer concretos, no puede refinar un AnyBuffer con Match: no hay un tipo de datos que se devuelve único. Aún se puede pasar a Match con un búfer ya escrito, cuyo tipo de datos que se devuelve ya se conoce. También se puede verificar un búfer ya escrito. Por lo general, esto es útil para verificar las relaciones de forma:

Error VerifySameShape(BufferR2<F32> lhs, BufferR2<F32> rhs) {
  int64_t rows;
  int64_t cols;

  Error error = Verify(
      "lhs", lhs,
      m::Buffer().WithDims(&rows, &cols));
  if (error.failure()) {
    return error;
  }

  return Verify("rhs", rhs,
                m::Buffer().WithDims(rows, cols));
}

El nombre que se pasa como primer argumento a Match y Verify es obligatorio y se incluye en los errores, junto con una descripción de la restricción fallida.

La API de FFI externa en xla/ffi/api/ffi.h muestra Error de Verify y ErrorOr<T> de Match, como en los ejemplos anteriores. La API de FFI interna en xla/ffi/ffi.h proporciona la misma interfaz de coincidencia con absl::Status y absl::StatusOr<T>:

ABSL_ASSIGN_OR_RETURN(BufferR2<F32> input,
                 Match("input", buffer, m::Buffer<F32, 2>()));
ABSL_RETURN_IF_ERROR(Verify(
    "input", input,
    m::Buffer().WithDims(expected_rows, m::Dim())));

Argumentos y resultados variádicos

Si la cantidad de argumentos y resultados puede ser diferente en diferentes instancias de una llamada personalizada, se pueden decodificar en el tiempo de ejecución con RemainingArgs y RemainingRets.

auto handler = Ffi::Bind().RemainingArgs().RemainingRets().To(
    [](RemainingArgs args, RemainingRets results) -> Error {
      ErrorOr<AnyBuffer> arg = args.get<AnyBuffer>(0);
      ErrorOr<Result<AnyBuffer>> res = results.get<AnyBuffer>(0);

      if (!arg.has_value()) {
        return Error(ErrorCode::kInternal, arg.error());
      }

      if (!res.has_value()) {
        return Error(ErrorCode::kInternal, res.error());
      }

      return Error::Success();
    });

Los argumentos y resultados variádicos se pueden declarar después de los argumentos y resultados normales. Sin embargo, vincular argumentos y resultados normales después de uno variádico es ilegal.

auto handler =
    Ffi::Bind()
        .Arg<AnyBuffer>()
        .RemainingArgs()
        .Ret<AnyBuffer>()
        .RemainingRets()
        .To([](AnyBuffer arg, RemainingArgs args, AnyBuffer ret,
               RemainingRets results) -> Error { return Error::Success(); });

Atributos

XLA FFI admite la decodificación automática de mlir::DictionaryAttr que se pasa como un backend_config custom_call en argumentos del controlador de FFI.

%0 = "stablehlo.custom_call"(%arg0) {
  call_target_name = "foo",
  backend_config= {
    i32 = 42 : i32,
    str = "string"
  },
  api_version = 4 : i32
} : (tensor<f32>) -> tensor<f32>

En este ejemplo, la llamada personalizada tiene un solo argumento de búfer y dos atributos, y XLA FFI puede decodificarlos automáticamente y pasarlos al elemento invocable definido por el usuario.

auto handler = Ffi::Bind()
  .Arg<BufferR0<F32>>()
  .Attr<int32_t>("i32")
  .Attr<std::string_view>("str")
  .To([](BufferR0<F32> buffer, int32_t i32, std::string_view str) {
    return Error::Success();
  });

Atributos de enumeración definidos por el usuario

XLA FFI puede decodificar automáticamente los atributos integrales de MLIR en enumeraciones definidas por el usuario. La clase de enumeración debe tener el mismo tipo integral subyacente, y la decodificación debe registrarse de forma explícita con XLA FFI.

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

enum class Command : int32_t {
  kAdd = 0,
  kMul = 1,
};

XLA_FFI_REGISTER_ENUM_ATTR_DECODING(Command);

auto handler = Ffi::Bind().Attr<Command>("command").To(
    [](Command command) -> Error { return Error::Success(); });

Cómo vincular todos los atributos de llamada personalizada

Es posible obtener acceso a todos los atributos de llamada personalizada como un diccionario y decodificar de forma diferida solo los atributos que se necesitan en el tiempo de ejecución.

auto handler = Ffi::Bind().Attrs().To([](Dictionary attrs) -> Error {
  ErrorOr<int32_t> i32 = attrs.get<int32_t>("i32");
  return Error::Success();
});

Atributos de estructura definidos por el usuario

XLA FFI puede decodificar atributos de diccionario en estructuras definidas por el usuario.

%0 = "stablehlo.custom_call"(%arg0) {
  call_target_name = "foo",
  backend_config= {
    range = { lo = 0 : i64, hi = 42 : i64 }
  },
  api_version = 4 : i32
} : (tensor<f32>) -> tensor<f32>

En el ejemplo anterior, range es un atributo mlir::DictionaryAttr y, en lugar de acceder a los campos del diccionario por nombre, se puede decodificar automáticamente como una estructura de C++. La decodificación debe registrarse de forma explícita con una macro XLA_FFI_REGISTER_STRUCT_ATTR_DECODING (en segundo plano, define una especialización de plantilla en el espacio de nombres ::xla::ffi, por lo que la macro se debe agregar al espacio de nombres global).

struct Range {
  int64_t lo;
  int64_t hi;
};

XLA_FFI_REGISTER_STRUCT_ATTR_DECODING(Range, StructMember<int64_t>("lo"),
                                             StructMember<int64_t>("hi"));

auto handler = Ffi::Bind().Attr<Range>("range").To([](Range range) -> Error{
  return Error::Success();
});

Los atributos personalizados se pueden cargar desde un diccionario, al igual que cualquier otro atributo. En el siguiente ejemplo, todos los atributos de llamada personalizada se decodifican como un Dictionary, y se puede acceder a un range por nombre.

auto handler = Ffi::Bind().Attrs().To([](Dictionary attrs) -> Error {
  ErrorOr<Range> range = attrs.get<Range>("range");
  return Error::Success();
});

Crea una llamada personalizada en la CPU

Puedes crear una instrucción HLO que represente una llamada personalizada a través de la API cliente de XLA. Por ejemplo, el siguiente código usa una llamada personalizada para calcular A[i] = B[i % 128]+ C[i] en la CPU. (Por supuesto, podrías, y deberías, hacer esto con HLO normal).

#include "xla/client/xla_builder.h"
#include "xla/service/custom_call_target_registry.h"

void do_it() {
  xla::XlaBuilder b("do_it");
  xla::XlaOp param0 =
      xla::Parameter(&b, 0, xla::ShapeUtil::MakeShape(xla::F32, {128}), "p0");
  xla::XlaOp param1 =
      xla::Parameter(&b, 1, xla::ShapeUtil::MakeShape(xla::F32, {2048}), "p1");
  xla::XlaOp custom_call =
      xla::CustomCall(&b, "do_custom_call", /*operands=*/{param0, param1},
        /*shape=*/xla::ShapeUtil::MakeShape(xla::F32, {2048}),
        /*opaque=*/"", /*has_side_effect=*/false,
        /*output_operand_aliasing=*/{}, /*literal=*/nullptr,
        /*schedule=*/CustomCallSchedule::SCHEDULE_NONE,
        /*api_version=*/CustomCallApiVersion::API_VERSION_TYPED_FFI);
}

// Constrain custom call arguments to 1-dimensional buffers of F32 data type.
using BufferF32 = xla::ffi::BufferR1<xla::ffi::DataType::F32>;

// Implement a custom call as a C++ function. Note that we can use `Buffer` type
// defined by XLA FFI that gives us access to buffer data type and shape.
xla::ffi::Error do_custom_call(BufferF32 in0, BufferF32 in1,
                               xla::ffi::Result<BufferF32> out) {
  size_t d0 = in0.dimensions[0];
  size_t d1 = in1.dimensions[0];

  // Check that dimensions are compatible.
  assert(out->dimensions[0] == d1 && "unexpected dimensions");

  for (size_t i = 0; i < d1; ++i) {
    out->data[i] = in0.data[i % d0] + in1.data[i];
  }
}

// Explicitly define an XLA FFI handler signature and bind it to the
// `do_custom_call` implementation. XLA FFI handler can automatically infer
// type signature from the custom call function, but it relies on magical
// template metaprogramming an explicit binding provides and extra level of
// type checking and clearly states custom call author intentions.
XLA_FFI_DEFINE_HANDLER(handler, do_custom_call,
                       ffi::Ffi::Bind()
                           .Arg<Buffer>()
                           .Arg<Buffer>()
                           .Ret<Buffer>());

// Registers `handler` with and XLA FFI on a "Host" platform.
XLA_FFI_REGISTER_HANDLER(xla::ffi::GetXlaFfiApi(), "do_custom_call",
                         "Host", handler);

Crea una llamada personalizada en la GPU

El registro de llamadas personalizadas de GPU con XLA FFI es casi idéntico. La única diferencia es que, para la GPU, debes solicitar una transmisión de plataforma subyacente (transmisión CUDA o ROCM) para poder iniciar el kernel en el dispositivo. A continuación, se muestra un ejemplo de CUDA que realiza el mismo cálculo (A[i] = B[i % 128] + C[i]) que el código de CPU anterior.

void do_it() { /* same implementation as above */ }

__global__ custom_call_kernel(const float* in0, const float* in1, float* out) {
  size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
  out[idx] = in0[idx % 128] + in1[idx];
}

void do_custom_call(CUstream stream, BufferF32 in0, BufferF32 in1,
                    xla::ffi::Result<BufferF32> out) {
  size_t d0 = in0.dimensions[0];
  size_t d1 = in1.dimensions[0];
  size_t d2 = out->dimensions[0];

  assert(d0 == 128 && d1 == 2048 && d2 == 2048 && "unexpected dimensions");

  const int64_t block_dim = 64;
  const int64_t grid_dim = 2048 / block_dim;
  custom_call_kernel<<<grid_dim, block_dim, 0, stream>>>(
    in0.data, in1.data, out->data);
}

XLA_FFI_DEFINE_HANDLER(handler, do_custom_call,
                       ffi::Ffi::Bind()
                           .Ctx<xla::ffi::PlatformStream<CUstream>>()
                           .Arg<BufferF32>()
                           .Arg<BufferF32>()
                           .Ret<BufferF32>());

XLA_FFI_REGISTER_HANDLER(xla::ffi::GetXlaFfiApi(), "do_custom_call",
                         "CUDA", handler);

Primero, ten en cuenta que la función de llamada personalizada de GPU sigue siendo una función que se ejecuta en la CPU. La función de CPU do_custom_call es responsable de poner en cola el trabajo en la GPU. Aquí, inicia un kernel de CUDA, pero también podría hacer otra cosa, como llamar a cuBLAS.

Los argumentos y los resultados también se encuentran en el host, y el miembro de datos contiene un puntero a la memoria del dispositivo (es decir, la GPU). Los búferes que se pasan al controlador de llamadas personalizadas tienen la forma de los búferes del dispositivo subyacente, por lo que la llamada personalizada puede calcular los parámetros de inicio del kernel a partir de ellos.

Cómo pasar tuplas a llamadas personalizadas

Considera la siguiente llamada personalizada.

using xla::ShapeUtil;
using xla::F32;
Shape p0_shape = ShapeUtil::MakeTuple({
    ShapeUtil::MakeShape(F32, {32}),
    ShapeUtil::MakeTuple({
        ShapeUtil::MakeShape(F32, {64}),
        ShapeUtil::MakeShape(F32, {128}),
    }),
    ShapeUtil::MakeShape(F32, {256}),
});
xla::XlaOp p0 = xla::Parameter(0, p0_shape, "p0");

Shape out_shape = ShapeUtil::MakeTuple({
  ShapeUtil::MakeShape(F32, {512}),
  ShapeUtil::MakeShape(F32, {1024}),
});
xla::CustomCall(&b, "do_custom_call", /*operands=*/{p0}, out_shape, ...);

Tanto en la CPU como en la GPU, una tupla se representa en la memoria como un array de punteros. Cuando XLA llama a llamadas personalizadas con argumentos o resultados de tupla, los aplana y los pasa como argumentos o resultados de búfer normales.

Salidas de tupla como búferes temporales

Las entradas de tupla a llamadas personalizadas son una comodidad, pero no son estrictamente necesarias. Si no admitiéramos entradas de tupla a llamadas personalizadas, siempre podrías descomprimir las tuplas con get-tuple-element antes de pasarlas a la llamada personalizada.

Por otro lado, las salidas de tupla te permiten hacer cosas que no podrías hacer de otra manera.

El motivo obvio para tener salidas de tupla es que las salidas de tupla son la forma en que una llamada personalizada (o cualquier otra operación de XLA) muestra varios arrays independientes.

Sin embargo, de forma menos obvia, una salida de tupla también es una forma de darle a tu llamada personalizada memoria temporal. Sí, una salida puede representar un búfer temporal. Considera que un búfer de salida tiene la propiedad de que la operación puede escribir en él y leerlo después de que se haya escrito. Eso es exactamente lo que quieres de un búfer temporal.

En el ejemplo anterior, supongamos que queríamos usar F32[1024] como un búfer temporal. Luego, escribiríamos el HLO como antes y nunca leeríamos el índice de tupla 1 de la salida de la llamada personalizada.