Chamadas personalizadas do XLA

Este documento descreve como escrever e usar chamadas personalizadas do XLA usando a biblioteca XLA FFI. A chamada personalizada é um mecanismo para descrever uma "operação" externa no módulo HLO para o compilador XLA (no tempo de compilação), e o XLA FFI é um mecanismo para registrar a implementação dessas operações com o XLA (no momento da execução). FFI significa "interface de função externa" e é um conjunto de APIs C que definem uma interface binária (ABI) para o XLA chamar um código externo escrito em outras linguagens de programação. O XLA fornece vinculações somente de cabeçalho para o XLA FFI escrito em C++, que oculta todos os detalhes de baixo nível das APIs C subjacentes do usuário final.

Chamadas personalizadas do JAX + XLA

Consulte a documentação do JAX para ver exemplos completos de integração de chamadas personalizadas e XLA FFI com o JAX.

Vinculação do XLA FFI

A vinculação do XLA FFI é uma especificação de tempo de compilação da assinatura de chamada personalizada: argumentos de chamada personalizada, atributos e tipos, além de outros parâmetros transmitidos pelo contexto de execução (ou seja, fluxo de GPU para back-end de GPU). A vinculação do XLA FFI pode ser vinculada a qualquer função C++ chamável (ponteiro de função, lambda etc.) com assinatura operator() compatível. O gerenciador construído decodifica o frame de chamada do XLA FFI (definido pela API C estável), verifica o tipo de todos os parâmetros e encaminha os resultados decodificados para o callback definido pelo usuário.

A vinculação do XLA FFI depende muito da metaprogramação de modelo para poder compilar o gerenciador construído no código de máquina mais eficiente. As sobrecargas de tempo de execução são da ordem de alguns nanossegundos para cada parâmetro de chamada personalizada.

Os pontos de personalização do XLA FFI são implementados como especializações de modelo, e os usuários podem definir como decodificar os tipos personalizados. Ou seja, é possível definir a decodificação personalizada para tipos enum class definidos pelo usuário.

Retornar erros de chamadas personalizadas

As implementações de chamadas personalizadas precisam retornar o valor xla::ffi::Error para sinalizar sucesso ou erro para o ambiente de execução do XLA. Ele é semelhante a absl::Status e tem o mesmo conjunto de códigos de erro. Não usamos absl::Status porque ele não tem uma ABI estável e seria inseguro transmiti-lo entre a biblioteca de chamadas personalizadas carregada dinamicamente e o próprio 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 e resultados de buffer

O XLA usa o estilo de transmissão de destino para resultados: chamadas personalizadas (ou qualquer outra operação do XLA) não alocam memória para resultados e, em vez disso, gravam em destinos transmitidos pelo ambiente de execução do XLA. O XLA usa a atribuição de buffer estático e aloca buffers para todos os valores com base nos intervalos ativos no tempo de compilação.

Os resultados transmitidos aos gerenciadores de FFI são encapsulados em um Result<T> modelo, que tem uma semântica semelhante a um ponteiro: operator-> dá acesso ao parâmetro subjacente.

Os argumentos e resultados AnyBuffer dão acesso aos parâmetros de buffer de chamada personalizada de qualquer tipo de dados. Isso é útil quando a chamada personalizada tem uma implementação genérica que funciona para vários tipos de dados, e a implementação de chamada personalizada faz o envio de tempo de execução com base no tipo de dados. AnyBuffer dá acesso ao tipo de dados do buffer, às dimensões e a um ponteiro para o próprio buffer.

%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 e resultados de buffer restritos

Buffer permite adicionar restrições ao tipo de dados do buffer e ao número de dimensões. Elas serão verificadas automaticamente pelo gerenciador e retornarão um erro para o ambiente de execução do XLA se os argumentos de tempo de execução não corresponderem à assinatura do gerenciador 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();
    });

Correspondência e verificação de buffers

Os padrões de buffer podem expressar restrições mais específicas do que os tipos de argumento e resultado em uma vinculação de FFI. Eles são úteis para refinar um AnyBuffer para um tipo de buffer concreto, verificar tamanhos de dimensão e verificar relações entre várias formas de buffer.

Os padrões são imutáveis e podem especificar dtype e classificação diretamente ou com os modificadores correspondentes:

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

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

A função Match verifica um padrão e retorna o tipo de buffer com restrição máxima. Para refinar um AnyBuffer, o padrão precisa especificar exatamente um tipo de dados e uma classificação. A correspondência de um buffer já digitado verifica outras restrições e preserva o 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();
    });

Neste exemplo, a correspondência de input captura os dois tamanhos de dimensão, e a correspondência de output verifica se ele tem o mesmo formato. As capturas só são confirmadas quando o padrão completo corresponde.

Um padrão de dimensão pode ser irrestrito, fixo ou capturado. Valores fixos e ponteiros de captura são convertidos implicitamente em padrões de dimensão:

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 descreve um formato completo e define a classificação com base no número de argumentos. WithDim<I> ou WithDim(index, ...) restringe as dimensões posicionalmente sem fixar a classificação completa. Uma restrição posicional exige que a dimensão referenciada exista, mas deixa todas as outras dimensões irrestritas:

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 exige um buffer F32 de classificação 4 com dimensões 0 e 3 iguais a 1 e 128, respectivamente. As dimensões 1 e 2 são irrestritas. Sem um explícito WithRank, WithDim(index, ...) aceita qualquer classificação maior que index.

As capturas de dimensão também podem expressar relações entre posições:

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

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

Quando o mesmo ponteiro de captura aparece várias vezes, todas as dimensões correspondentes precisam ter o mesmo tamanho. Portanto, a chamada Match acima aceita apenas matrizes quadradas. As capturas de dimensão são gravadas somente após o padrão completo ser bem-sucedido e não são alteradas em caso de falha.

As relações de buffer inteiro não exigem a captura de todas as dimensões. WithShapeOf corresponde ao formato de tempo de execução completo de outro buffer, permitindo um dtype diferente. Like também exige o mesmo 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));
}

Os dois modificadores copiam os metadados do buffer de referência quando o padrão é criado. O padrão não mantém uma referência ao buffer. Como essas são restrições de tempo de execução, elas não refinam um AnyBuffer para um tipo de retorno concreto. Elas são úteis principalmente com Verify ou com Match quando o buffer de entrada já tem um tipo concreto.

Use Verify quando o buffer precisar manter o tipo atual ou quando um padrão aceitar mais de um tipo de dados ou classificação. Por exemplo, um buffer de índice pode aceitar S32 ou S64 e enviar no tipo de dados após a verificação:

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

Como esse padrão permite vários tipos de buffer concretos, ele não pode refinar um AnyBuffer com Match: não há um tipo de retorno exclusivo. Ele ainda pode ser transmitido para Match com um buffer já digitado, cujo tipo de retorno já é conhecido. Um buffer já digitado também pode ser verificado. Isso geralmente é útil para verificar relações de formato:

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

O nome transmitido como o primeiro argumento para Match e Verify é obrigatório e incluído em erros, juntamente com uma descrição da restrição com falha.

A API FFI externa em xla/ffi/api/ffi.h retorna Error de Verify e ErrorOr<T> de Match, como nos exemplos acima. A API FFI interna em xla/ffi/ffi.h fornece a mesma interface correspondente usando absl::Status e 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 e resultados variádicos

Se o número de argumentos e resultados puder ser diferente em instâncias diferentes de uma chamada personalizada, eles poderão ser decodificados no momento da execução usando RemainingArgs e 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();
    });

Argumentos e resultados variádicos podem ser declarados após argumentos e resultados normais. No entanto, a vinculação de argumentos e resultados normais após um variádico é 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

O XLA FFI oferece suporte à decodificação automática de mlir::DictionaryAttr transmitido como um backend_config custom_call em argumentos do gerenciador 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>

Neste exemplo, a chamada personalizada tem um único argumento de buffer e dois atributos, e o XLA FFI pode decodificá-los automaticamente e transmitir para a função chamável definida pelo usuário.

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 enum definidos pelo usuário

O XLA FFI pode decodificar automaticamente atributos MLIR integrais em enums definidos pelo usuário. A classe enum precisa ter o mesmo tipo integral subjacente, e a decodificação precisa ser registrada explicitamente com o 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(); });

Vincular todos os atributos de chamada personalizada

É possível acessar todos os atributos de chamada personalizada como um dicionário e decodificar lentamente apenas os atributos necessários no momento da execução.

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

Atributos de struct definidos pelo usuário

O XLA FFI pode decodificar atributos de dicionário em structs definidos pelo usuário.

%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>

No exemplo acima, range é um atributo mlir::DictionaryAttr e, em vez de acessar campos de dicionário por nome, ele pode ser decodificado automaticamente como um struct C++. A decodificação precisa ser registrada explicitamente com uma macro XLA_FFI_REGISTER_STRUCT_ATTR_DECODING (nos bastidores, ela define uma especialização de modelo no namespace ::xla::ffi. Portanto, a macro precisa ser adicionada ao namespace 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();
});

Os atributos personalizados podem ser carregados de um dicionário, assim como qualquer outro atributo. No exemplo abaixo, todos os atributos de chamada personalizada são decodificados como um Dictionary, e um range pode ser acessado por nome.

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

Criar uma chamada personalizada na CPU

É possível criar uma instrução HLO que representa uma chamada personalizada pela API de cliente do XLA. Por exemplo, o código a seguir usa uma chamada personalizada para calcular A[i] = B[i % 128]+ C[i] na CPU. É claro que você pode (e deve! ) fazer isso com o 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);

Criar uma chamada personalizada na GPU

O registro de chamada personalizada da GPU com o XLA FFI é quase idêntico. A única diferença é que, para a GPU, é necessário pedir um fluxo de plataforma subjacente (fluxo CUDA ou ROCM) para poder iniciar o kernel no dispositivo. Confira um exemplo de CUDA que faz o mesmo cálculo (A[i] = B[i % 128] + C[i]) que o código da CPU acima.

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);

Primeiro, observe que a função de chamada personalizada da GPU ainda é uma função executada na CPU. A função de CPU do_custom_call é responsável por enfileirar o trabalho na GPU. Aqui, ela inicia um kernel CUDA, mas também pode fazer outra coisa, como chamar cuBLAS.

Os argumentos e resultados também ficam no host, e o membro de dados contém um ponteiro para a memória do dispositivo (ou seja, GPU). Os buffers transmitidos ao gerenciador de chamadas personalizadas têm o formato dos buffers de dispositivo subjacentes. Assim, a chamada personalizada pode calcular os parâmetros de inicialização do kernel.

Transmitir tuplas para chamadas personalizadas

Considere a seguinte chamada 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, ...);

Na CPU e na GPU, uma tupla é representada na memória como uma matriz de ponteiros. Quando o XLA chama chamadas personalizadas com argumentos ou resultados de tupla, ele os achata e transmite como argumentos ou resultados de buffer normais.

Saídas de tupla como buffers temporários

As entradas de tupla para chamadas personalizadas são convenientes, mas não são estritamente necessárias. Se não oferecêssemos suporte a entradas de tupla para chamadas personalizadas, você sempre poderia descompactar as tuplas usando get-tuple-element antes de transmiti-las para a chamada personalizada.

Por outro lado, as saídas de tupla permitem fazer coisas que você não poderia fazer de outra forma.

O motivo óbvio para ter saídas de tupla é que elas são a maneira como uma chamada personalizada (ou qualquer outra operação do XLA) retorna várias matrizes independentes.

Mas, menos obviamente, uma saída de tupla também é uma maneira de dar memória temporária à chamada personalizada. Sim, uma saída pode representar um buffer temporário. Considere que um buffer de saída tem a propriedade de que a operação pode gravar nele e pode ler dele depois que ele for gravado. É exatamente isso que você quer de um buffer temporário.

No exemplo acima, suponha que queremos usar o F32[1024] como um buffer temporário. Em seguida, escreveríamos o HLO como acima e simplesmente nunca leríamos o índice de tupla 1 da saída da chamada personalizada.