Questa pagina è stata tradotta dall'API Cloud Translation.
Switch to English

Chiamate personalizzate XLA

Questo documento descrive come scrivere e utilizzare le "chiamate personalizzate" XLA. Le chiamate personalizzate consentono di richiamare il codice scritto in un linguaggio di programmazione come C ++ o CUDA da un programma XLA.

Chiamata personalizzata su CPU

È possibile creare un'istruzione HLO che rappresenta una chiamata personalizzata tramite l'API client XLA. Questo non è esposto tramite TensorFlow al momento della scrittura.

Ad esempio, il codice seguente utilizza una chiamata personalizzata per calcolare A[i] = B[i % 128] + C[i] sulla CPU. (Certo che potresti - e dovresti! - farlo con un normale HLO.)

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

Si noti che la funzione do_custom_call deve conoscere le dimensioni dei buffer su cui opera. In questo esempio codifichiamo le dimensioni 128 e 2048. Se non si desidera farlo, è possibile passare le dimensioni come parametri alla chiamata.

Chiamata personalizzata su GPU

Il framework di chiamate personalizzate GPU è leggermente diverso da quello sulla CPU. Ecco un esempio CUDA che esegue lo stesso calcolo A[i] = B[i % 128] + C[i] del codice CPU sopra.

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

Si noti innanzitutto che la funzione di chiamata personalizzata GPU è ancora una funzione eseguita sulla CPU . La nostra funzione CPU do_custom_call è responsabile del lavoro di accodamento sulla GPU. Qui avvia un kernel CUDA, ma potrebbe anche fare qualcos'altro, come call cublas.

buffers è un array di puntatori che risiede sull'host e ogni elemento che contiene punta alla memoria del dispositivo (ad es. GPU). I parametri vengono prima, seguiti dal valore di uscita. Ciò è notevolmente diverso dalla convenzione di chiamata della CPU, che ha due parametri, ins e out . Il motivo principale per cui divergiamo è rendere possibile la gestione efficiente di input / output a forma di tupla; vedere la sezione seguente.

Come nell'esempio della CPU, abbiamo inserito le dimensioni del buffer di input e output nella nostra chiamata personalizzata. Tuttavia, diversamente dal caso della CPU, passare le dimensioni del buffer come operandi alla chiamata personalizzata non funzionerebbe bene. Di solito abbiamo bisogno delle dimensioni del buffer disponibili sulla CPU; ad esempio quando si avvia un kernel, è necessario conoscere le dimensioni di blocco / griglia da utilizzare. Ma se dovessimo passare le dimensioni del buffer come operandi alla nostra chiamata personalizzata, i loro valori vivrebbero nella memoria della GPU. Dovremmo quindi fare una costosa memoria sincrona da dispositivo a host all'inizio della nostra operazione solo per leggere le dimensioni.

Per consentirti di aggirare questo problema, forniamo il parametro opaque . È possibile impostare questo su una stringa arbitraria di byte quando si crea la chiamata personalizzata:

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

Poiché xla::Shape ha una rappresentazione di buffer di protocollo, è possibile memorizzare questo proto serializzato all'interno di opaque e deserializzarlo nella chiamata personalizzata della GPU. Si noti tuttavia che, sebbene xla::ShapeProto non cambia frequentemente, lo fa cambiare. Controlla il registro git per vedere come è cambiato in passato.

Passando le tuple alle chiamate personalizzate

Considera la seguente chiamata personalizzata.

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

Sia sulla CPU che sulla GPU, una tupla è rappresentata in memoria come una matrice di puntatori. In C ++ - pseudocodice, il parametro 0 sopra è presentato come segue.

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

Sebbene la rappresentazione in memoria delle tuple sia la stessa in CPU e GPU, sono gestite in modo diverso nelle convenzioni di chiamata su chiamata personalizzata CPU e GPU.

Tuple output come buffer di temperatura

Gli input di tupla alle chiamate personalizzate sono una comodità, ma non sono strettamente necessari. Se non supportiamo gli input di tupla alle chiamate personalizzate, è sempre possibile decomprimere le tuple utilizzando get-tuple-element prima di passarle alla chiamata personalizzata.

D'altra parte, le uscite a tupla ti consentono di fare cose che altrimenti non potresti fare.

Il motivo ovvio per avere output di tupla è che è così che una chiamata personalizzata (o qualsiasi altra operazione XLA) restituisce più array indipendenti.

Ma meno ovviamente, un output di tupla è anche un modo per dare alla tua memoria di chiamata temporanea personalizzata. Sì, un output può rappresentare un buffer temporaneo. Considera, un buffer di output ha la proprietà che l'op può scrivere su di esso e può leggerlo dopo che è stato scritto. Questo è esattamente quello che vuoi da un buffer temporaneo.

Nell'esempio sopra, supponiamo di voler usare l' F32[1024] come buffer temporaneo. Quindi scriveremmo l'HLO come sopra e semplicemente non leggeremmo mai l'indice tupla 1 dell'output della chiamata personalizzata.

Tuple nelle chiamate personalizzate della CPU

Nel codice CPU, abbiamo una funzione do_custom_call(const void** ins, void* out) . ins è un array con un solo elemento, che punta a param0 . I subbuffers di param0 sono accessibili da dereferenziazione che puntatore, e le subbuffers di output_tuple sono accessibili da dereferenziazione out .

Tuple nelle chiamate personalizzate GPU

Nel codice GPU, abbiamo una funzione do_custom_call(..., void** buffers, ...) . In questo caso buffers è un array host di sei puntatori di dispositivo, uno per ciascun buffer foglia nell'input / output. Per generare l'elenco semplice, ripetiamo i parametri e l'output e per ciascuno di noi eseguiamo un attraversamento del preordine della sua forma. In concreto:

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