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 codice scritto in un linguaggio di programmazione come C ++ o CUDA da un programma XLA.

Chiamata personalizzata sulla CPU

È possibile creare un'istruzione HLO che rappresenta una chiamata personalizzata tramite l'API client di 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. (Ovviamente potresti - e dovresti! - farlo con il 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 vuoi farlo, puoi passare le dimensioni come parametri alla chiamata.

Chiamata personalizzata su GPU

Il framework delle chiamate personalizzate della GPU è leggermente diverso da quello della 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 della GPU è ancora una funzione eseguita sulla CPU . La nostra funzione CPU do_custom_call è responsabile del lavoro di accodamento sulla GPU. Qui lancia un kernel CUDA, ma potrebbe anche fare qualcos'altro, come chiamare cublas.

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

Come nell'esempio della CPU, abbiamo codificato le dimensioni del buffer di input e output nella nostra chiamata personalizzata. Tuttavia, a differenza del caso della CPU, passare le dimensioni del buffer come operandi alla chiamata personalizzata non funzionerebbe bene. Di solito abbiamo bisogno delle dimensioni del buffer a nostra disposizione sulla CPU; ad esempio, quando si avvia un kernel, è necessario conoscere le dimensioni del 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 un costoso memcpy sincrono da dispositivo a host all'inizio della nostra operazione solo per leggere le dimensioni.

Per consentirti di aggirare questo problema, forniamo il parametro opaque . Puoi impostarlo su una stringa arbitraria di byte quando crei 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 del buffer del protocollo, è possibile memorizzare questo proto serializzato all'interno di opaque e deserializzarlo all'interno della chiamata personalizzata della GPU. Si noti tuttavia che, sebbene xla::ShapeProto non cambia frequentemente, lo fa cambiare. Controlla il registro di git per vedere come è cambiato in passato.

Passaggio di 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 un array di puntatori. Nello pseudocodice C ++, il parametro 0 sopra è disposto 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, vengono gestite in modo diverso nelle convenzioni di chiamata personalizzate di CPU e GPU.

Tupla output come buffer temporanei

Gli input di tuple per le chiamate personalizzate sono una comodità, ma non sono strettamente necessari. Se non supportassimo gli input delle tuple per le chiamate personalizzate, potresti sempre decomprimere le tuple usando get-tuple-element prima di passarle alla chiamata personalizzata.

D'altra parte, gli output delle tuple ti consentono di fare cose che altrimenti non potresti.

L'ovvia ragione per avere output di tuple è che è così che una chiamata personalizzata (o qualsiasi altra operazione XLA) restituisce più array indipendenti.

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

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

Tuple nelle chiamate personalizzate della CPU

Nel codice della CPU, abbiamo una funzione do_custom_call(const void** ins, void* out) . ins è un array con un solo elemento, che punta a param0 . I subuffer di param0 sono accessibili dereferenziando quel puntatore, e i subuffer di output_tuple sono accessibili dereferenziando out .

Tuple nelle chiamate personalizzate GPU

Nel codice GPU, abbiamo una funzione do_custom_call(..., void** buffers, ...) . In questo caso, i buffers sono un array host di sei puntatori di dispositivo, uno per ogni buffer foglia nell'input / output. Per generare l'elenco semplice, iteriamo sui parametri e sull'output e per ognuno 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