Ta strona została przetłumaczona przez Cloud Translation API.
Switch to English

Połączenia niestandardowe XLA

W tym dokumencie opisano, jak pisać i używać niestandardowych wywołań XLA. Wywołania niestandardowe umożliwiają wywoływanie kodu napisanego w języku programowania, takim jak C ++ lub CUDA, z programu XLA.

Niestandardowe wywołanie procesora

Możesz utworzyć instrukcję HLO, która reprezentuje wywołanie niestandardowe za pośrednictwem interfejsu API klienta XLA. Nie jest to ujawniane przez TensorFlow w momencie pisania.

Na przykład poniższy kod wykorzystuje wywołanie niestandardowe do obliczenia A[i] = B[i % 128] + C[i] na procesorze. (Oczywiście, że możesz - i powinieneś! - to zrobić za pomocą zwykłego 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");
 

Zauważ, że funkcja do_custom_call musi znać wymiary buforów, na których działa. W tym przykładzie zakodowaliśmy na stałe rozmiary 128 i 2048. Jeśli nie chcesz tego robić, możesz przekazać wymiary jako parametry do wywołania.

Niestandardowe wywołanie GPU

Struktura niestandardowych wywołań GPU różni się nieco od struktury procesora. Oto przykład CUDA, który wykonuje te same obliczenia A[i] = B[i % 128] + C[i] co powyższy kod procesora.

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

Zauważ najpierw, że funkcja niestandardowego wywołania GPU jest nadal funkcją wykonywaną na procesorze . Nasza do_custom_call procesora do_custom_call jest odpowiedzialna za kolejkowanie pracy na GPU. Tutaj uruchamia jądro CUDA, ale może też zrobić coś innego, na przykład call cublas.

buffers to tablica wskaźników, która żyje na hoście, a każdy element, który zawiera, wskazuje na pamięć urządzenia (np. GPU). Najpierw parametry, po których następuje wartość wyjściowa. Różni się to w szczególności od konwencji wywoływania procesora, która ma dwa parametry, ins i out . Głównym powodem, dla którego się rozchodzimy, jest umożliwienie wydajnej obsługi wejść / wyjść w kształcie krotki; zobacz sekcję poniżej.

Podobnie jak w przykładzie z procesorem, zakodowaliśmy na stałe rozmiary buforów wejściowych i wyjściowych w naszym wywołaniu niestandardowym. Jednak inaczej niż w przypadku procesora, przekazywanie rozmiarów buforów jako operandy do wywołania niestandardowego nie działałoby dobrze. Zwykle potrzebujemy dostępnych rozmiarów buforów na procesorze; np. uruchamiając jądro, musimy znać wymiary bloku / siatki, których będziemy używać. Ale gdybyśmy przekazali rozmiary buforów jako operandy do naszego niestandardowego wywołania, ich wartości znajdowałyby się w pamięci GPU. Na początku naszej operacji musielibyśmy przeprowadzić kosztowną synchroniczną memcpy typu urządzenie-host, aby odczytać rozmiary.

Abyś mógł obejść ten problem, podajemy opaque parametr. Możesz ustawić to na dowolny ciąg bajtów podczas tworzenia niestandardowego wywołania:

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

Ponieważ xla::Shape ma reprezentację bufora protokołu, możesz przechowywać ten serializowany proto wewnątrz opaque i deserializować go w niestandardowym wywołaniu GPU. Należy jednak pamiętać, że chociaż xla::ShapeProto nie zmienia się często, to się zmienia. Sprawdź dziennik gita, aby zobaczyć, jak zmienił się w przeszłości.

Przekazywanie krotek do wywołań niestandardowych

Rozważ następujące wywołanie niestandardowe.

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

Zarówno na CPU, jak i GPU, krotka jest reprezentowana w pamięci jako tablica wskaźników. W C ++ - pseudokodzie parametr 0 powyżej jest przedstawiony w następujący sposób.

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

Chociaż reprezentacja krotek w pamięci jest taka sama w CPU i GPU, są one obsługiwane inaczej w konwencjach wywoływania niestandardowych wywołań procesora i GPU.

Wyjścia krotki jako bufory tymczasowe

Dane wejściowe krotki do niestandardowych wywołań są wygodą, ale nie są bezwzględnie konieczne. Jeśli nie obsługujemy danych wejściowych krotki w wywołaniach niestandardowych, zawsze można rozpakować krotki za pomocą elementu get-tuple-element przed przekazaniem ich do wywołania niestandardowego.

Z drugiej strony dane wyjściowe krotki pozwalają robić rzeczy, których inaczej nie można byłoby zrobić.

Oczywistym powodem posiadania danych wyjściowych krotki jest to, że wywołanie niestandardowe (lub jakakolwiek inna operacja XLA) zwraca wiele niezależnych tablic.

Ale co mniej oczywiste, wyjście krotki jest również sposobem na udostępnienie niestandardowej pamięci tymczasowej wywołania. Tak, dane wyjściowe mogą reprezentować tymczasowy bufor. Rozważmy, że bufor wyjściowy ma właściwość, którą op może do niego zapisać, i może z niego odczytać po zapisaniu. Dokładnie tego chcesz od bufora tymczasowego.

W powyższym przykładzie załóżmy, że chcemy użyć F32[1024] jako bufora tymczasowego. Następnie zapisywalibyśmy HLO tak jak powyżej i po prostu nigdy nie czytalibyśmy indeksu krotki 1 wyniku niestandardowego wywołania.

Krotki w wywołaniach niestandardowych procesora

W kodzie procesora mamy funkcję do_custom_call(const void** ins, void* out) . ins to tablica zawierająca tylko jeden element, który wskazuje na param0 . W subbuffers z param0 są dostępne dereferencing tego wskaźnika, a subbuffers z output_tuple są dostępne dereferencing out .

Krotki w niestandardowych wywołaniach GPU

W kodzie GPU mamy funkcję do_custom_call(..., void** buffers, ...) . W tym przypadku buffers to tablica hostów składająca się z sześciu wskaźników urządzeń, po jednym dla każdego bufora liścia na wejściu / wyjściu. Aby wygenerować płaską listę, iterujemy po parametrach i danych wyjściowych, a dla każdego wykonujemy przeglądanie jego kształtu przed zamówieniem. Konkretnie:

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