Esta página foi traduzida pela API Cloud Translation.
Switch to English

Chamadas personalizadas XLA

Este documento descreve como escrever e usar XLA "chamadas personalizadas". As chamadas personalizadas permitem invocar o código escrito em uma linguagem de programação como C ++ ou CUDA a partir de um programa XLA.

Chamada personalizada na CPU

Você pode criar uma instrução HLO que represente uma chamada personalizada via API do cliente do XLA. Isso não é exposto através do TensorFlow no momento da gravação.

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 regularmente.)

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

void do_it() {
  xla::XlaBuilder b("do_it");
  xla::XlaOp param0 =
      xla::Parameter(0, xla::ShapeUtil::CreateShape(F32, {128}), "p0");
  xla::XlaOp param1 =
      xla::Parameter(1, xla::ShapeUtil::CreateShape(F32, {2048}), "p1");
  xla::XlaOp custom_call =
      xla::CustomCall(&b, "do_custom_call", /*operands=*/{param0, param1},
                      /*output_shape=*/ShapeUtil::CreateShape(F32, {2048}));
}

void do_custom_call(void* out, const void** in) {
  float* out_buf = reinterpret_cast<float*>(out);
  const float* in0 = reinterpret_cast<const float*>(in[0]);
  const float* in1 = reinterpret_cast<const float*>(in[1]);
  for (int i = 0; i < 2048; ++i) {
    out_buf[i] = in0[i % 128] + in1[i];
  }
}
XLA_REGISTER_CUSTOM_CALL_TARGET(do_custom_call, "Host");
 

Observe que a função do_custom_call precisa conhecer as dimensões dos buffers nos quais opera. Neste exemplo, codificamos os tamanhos 128 e 2048. Se você não quiser fazer isso, pode passar as dimensões como parâmetros para a chamada.

Chamada personalizada na GPU

A estrutura de chamada personalizada da GPU é um pouco diferente da da CPU. Aqui está um exemplo 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 = threadIdx.x * blockSize.x + gridIdx.x;
  out[idx] = in0[idx % 128] + in1[idx];
}

void do_custom_call(CUstream stream, void** buffers,
                    const char* opaque, size_t opaque_len) {
  const float* in0 = reinterpret_cast<const float*>(buffers[0]);
  const float* in1 = reinterpret_cast<const float*>(buffers[1]);
  float* out = reinterpret_cast<float*>(buffers[2]);

  const int64 block_dim = 64;
  const int64 grid_dim = 2048 / block_dim;
  custom_call_kernel<<<grid_dim, block_dim,
                       /*dynamic_shared_mem_bytes=*/0, stream>>>(in0, in1, out);
}
XLA_REGISTER_CUSTOM_CALL_TARGET(do_custom_call, "CUDA");
 

Observe primeiro que a função de chamada personalizada da GPU ainda é uma função executada na CPU . Nossa função da CPU do_custom_call é responsável pelo enfileiramento do trabalho na GPU. Aqui, ele lança um kernel CUDA, mas também pode fazer outra coisa, como chamar cublas.

buffers é uma matriz de ponteiros que fica no host e cada elemento contém pontos para a memória do dispositivo (por exemplo, GPU). Os parâmetros vêm primeiro, seguidos pelo valor de saída. Isso é notavelmente diferente da convenção de chamada da CPU, que possui dois parâmetros, ins e out . A principal razão pela qual divergimos é possibilitar o tratamento eficiente de entradas / saídas em forma de tupla; veja a seção abaixo.

Como no exemplo da CPU, codificamos os tamanhos do buffer de entrada e saída em nossa chamada personalizada. No entanto, diferentemente do caso da CPU, passar os tamanhos do buffer como operandos para a chamada personalizada não funcionaria bem. Normalmente, precisamos dos tamanhos de buffer disponíveis para nós na CPU; por exemplo, ao iniciar um kernel, precisamos conhecer as dimensões do bloco / grade a serem usadas. Porém, se passarmos os tamanhos de buffer como operandos para nossa chamada personalizada, seus valores permanecerão na memória da GPU. Em seguida, teríamos que fazer uma memória síncrona cara de dispositivo para host no início de nossa operação apenas para ler os tamanhos.

Para que você possa solucionar isso, fornecemos o parâmetro opaque . Você pode definir isso como uma sequência arbitrária de bytes ao criar a chamada personalizada:

 std::string opaque = "...";
xla::CustomCall(&b, "do_custom_call", /*operands=*/{param0, param1},
                /*output_shape=*/ShapeUtil::CreateShape(F32, {2048}),
                opaque);
 

Como o xla::Shape possui uma representação de buffer de protocolo, você pode armazenar esse proto serializado dentro do opaque e desserializar na sua chamada personalizada da GPU. Note, porém, que, embora xla::ShapeProto não muda com freqüência, ele faz a mudança. Verifique o log do git para ver como ele mudou no passado.

Passando tuplas para chamadas personalizadas

Considere a seguinte chamada personalizada.

 using xla::ShapeUtil;
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. No pseudocódigo C ++, o parâmetro 0 acima é apresentado da seguinte maneira.

 // In-memory layout of parameter 0 from custom-call above.  True on both CPU
// and GPU.
float* subbuf0 = new float[32];
float* subbuf1 = new float[64];
float* subbuf2 = new float[128]
float* subbuf3 = new float[256];

void* subtuple = new void*[2];
(*subtuple)[0] = subbuf1;
(*subtuple)[1] = subbuf2;

void* p0 = new void*[3];
(*p0)[0] = subbuf0;
(*p0)[1] = subtuple;
(*p0)[2] = subbuf3;
 

Embora a representação na memória das tuplas seja a mesma na CPU e na GPU, elas são tratadas de maneira diferente nas convenções de chamada de chamada personalizada da CPU e da GPU.

Saídas tuplas como buffers temporários

Entradas tuplas para chamadas personalizadas são uma conveniência, mas não são estritamente necessárias. Se não suportássemos entradas de tupla para chamadas personalizadas, você sempre pode descompactar as tuplas usando o elemento get-tuple antes de passá-las para a chamada personalizada.

Por outro lado, as saídas da tupla permitem que você faça coisas que não poderia de outra maneira.

A razão óbvia para obter saídas de tupla é que é assim que uma chamada personalizada (ou qualquer outra operação XLA) retorna várias matrizes independentes.

Porém, menos obviamente, uma saída de tupla também é uma maneira de fornecer à memória temporária de chamadas personalizada. Sim, uma saída pode representar um buffer temporário. Considere, um buffer de saída tem a propriedade que o op pode gravar nele, e pode ler dele depois de ter sido gravado. É exatamente isso que você deseja de um buffer temporário.

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

Tuplas em chamadas personalizadas da CPU

No código da CPU, temos a função do_custom_call(const void** ins, void* out) . ins é uma matriz com apenas um elemento, que aponta para param0 . Os subbuffers de param0 são acessíveis desreferenciando esse ponteiro, e os subbuffers de output_tuple são acessíveis desreferenciando a out .

Tuplas em chamadas personalizadas da GPU

No código da GPU, temos a função do_custom_call(..., void** buffers, ...) . Nesse caso, os buffers são uma matriz de host de seis ponteiros de dispositivo, um para cada buffer de folha na entrada / saída. Para gerar a lista simples, iteramos sobre os parâmetros e a saída e, para cada um, fazemos um percurso de pré-encomenda de sua forma. Concretamente:

 // Layout of `buffers` parameter to GPU custom call function for custom-call
// above.
buffers[0] == subbuf0
buffers[1] == subbuf1
buffers[2] == subbuf2
buffers[3] == subbuf3
buffers[4] == output_subbuf0
buffers[5] == output_subbuf1