Se usó la API de Cloud Translation para traducir esta página.
Switch to English

Llamadas personalizadas de XLA

Este documento describe cómo escribir y usar XLA "llamadas personalizadas". Las llamadas personalizadas le permiten invocar código escrito en un lenguaje de programación como C ++ o CUDA desde un programa XLA.

Llamada personalizada en la CPU

Puede crear una instrucción HLO que represente una llamada personalizada a través de la API de cliente de XLA. Esto no se expone a través de TensorFlow al momento de la escritura.

Por ejemplo, el siguiente código utiliza una llamada personalizada para calcular A[i] = B[i % 128] + C[i] en la CPU. (¡Por supuesto que podrías, y deberías!) Hacer esto con HLO normal).

 #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 la función do_custom_call necesita conocer las dimensiones de los buffers sobre los que opera. En este ejemplo, codificamos los tamaños 128 y 2048. Si no desea hacer esto, puede pasar las dimensiones como parámetros a la llamada.

Llamada personalizada en GPU

El marco de llamadas personalizado de GPU es algo diferente al de la CPU. Aquí hay un ejemplo de CUDA que hace 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 = 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 primero que la función de llamada personalizada de GPU sigue siendo una función ejecutada en la CPU . Nuestra función de CPU do_custom_call es responsable de poner en cola el trabajo en la GPU. Aquí lanza un núcleo CUDA, pero también podría hacer algo más, como llamar a cublas.

buffers es una matriz de punteros que vive en el host, y cada elemento que contiene apunta a la memoria del dispositivo (es decir, GPU). Los parámetros vienen primero, seguidos del valor de salida. Esto es notablemente diferente de la convención de llamadas de CPU, que tiene dos parámetros, ins y out . La razón principal por la que divergemos es para permitir el manejo eficiente de las entradas / salidas en forma de tupla; Vea la sección a continuación.

Como en el ejemplo de la CPU, hemos codificado los tamaños del búfer de entrada y salida en nuestra llamada personalizada. Sin embargo, a diferencia del caso de la CPU, pasar los tamaños del búfer como operandos a la llamada personalizada no funcionaría bien. Por lo general, necesitamos los tamaños de búfer disponibles en la CPU; Por ejemplo, al iniciar un kernel, necesitamos conocer las dimensiones de bloque / cuadrícula para usar. Pero si tuviéramos que pasar los tamaños del búfer como operandos a nuestra llamada personalizada, sus valores vivirían en la memoria de la GPU. Luego tendríamos que hacer una costosa memoria síncrona de dispositivo a host al comienzo de nuestra operación solo para leer los tamaños.

Para que pueda evitar esto, proporcionamos el parámetro opaque . Puede establecer esto en una cadena arbitraria de bytes cuando cree la llamada personalizada:

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

Dado que xla::Shape tiene una representación de búfer de protocolo, puede almacenar este prototipo serializado dentro de opaque y deserializarlo dentro de su llamada personalizada de GPU. Sin embargo, xla::ShapeProto cuenta que aunque xla::ShapeProto no cambia con frecuencia, cambia. Consulte el registro de git para ver cómo ha cambiado en el pasado.

Pasar tuplas a llamadas personalizadas

Considere la siguiente llamada 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);
 

Tanto en la CPU como en la GPU, una tupla se representa en la memoria como una matriz de punteros. En C ++ - pseudocódigo, el parámetro 0 anterior se presenta de la siguiente manera.

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

Aunque la representación en memoria de las tuplas es la misma en CPU y GPU, se manejan de manera diferente en las convenciones de llamadas de llamadas personalizadas de CPU y GPU.

Las salidas de tupla como amortiguadores temporales

Las entradas de tupla para llamadas personalizadas son convenientes, pero no son estrictamente necesarias. Si no admitimos entradas de tuplas para llamadas personalizadas, siempre puede desempaquetar las tuplas usando get-tuple-element antes de pasarlas a la llamada personalizada.

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

La razón obvia para tener salidas de tupla es que así es como una llamada personalizada (o cualquier otra operación XLA) devuelve múltiples arreglos independientes.

Pero menos obvio, una salida de tupla también es una forma de darle a su memoria temporal de llamada personalizada. Sí, una salida puede representar un búfer temporal. Tenga en cuenta que un búfer de salida tiene la propiedad de que el operador puede escribir en él y puede leerlo después de haberlo escrito. Eso es exactamente lo que quieres de un buffer temporal.

En el ejemplo anterior, supongamos que queremos usar el F32[1024] como un buffer temporal. Luego, escribiríamos el HLO tal como se indicó anteriormente, y simplemente nunca leeríamos el índice 1 de tupla de la salida de la llamada personalizada.

Tuplas en llamadas personalizadas de CPU

En el código de la CPU, tenemos una función do_custom_call(const void** ins, void* out) . ins es una matriz con solo un elemento, que apunta a param0 . Los subbuffers de param0 son accesibles por eliminación de referencias a ese puntero, y los subbuffers de output_tuple son accesibles mediante la eliminación de referencias out .

Tuplas en llamadas personalizadas de GPU

En el código GPU, tenemos una función do_custom_call(..., void** buffers, ...) . En este caso, los buffers son una matriz de host de seis punteros de dispositivo, uno para cada buffer de hoja en la entrada / salida. Para generar la lista plana, iteramos sobre los parámetros y la salida, y para cada uno realizamos un recorrido previo de su 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