Diese Seite wurde von der Cloud Translation API übersetzt.
Switch to English

Benutzerdefinierte XLA-Anrufe

In diesem Dokument wird beschrieben, wie Sie benutzerdefinierte XLA-Aufrufe schreiben und verwenden. Mit benutzerdefinierten Aufrufen können Sie Code aufrufen, der in einer Programmiersprache wie C ++ oder CUDA aus einem XLA-Programm geschrieben wurde.

Benutzerdefinierter Aufruf der CPU

Sie können eine HLO-Anweisung erstellen, die einen benutzerdefinierten Aufruf über die Client-API von XLA darstellt. Dies wird zum Zeitpunkt des Schreibens nicht über TensorFlow angezeigt.

Der folgende Code verwendet beispielsweise einen benutzerdefinierten Aufruf, um A[i] = B[i % 128] + C[i] auf der CPU zu berechnen. (Natürlich können - und sollten! - Sie dies mit normalem HLO tun.)

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

Beachten Sie, dass die Funktion do_custom_call die Abmessungen der Puffer kennen muss, über die sie arbeitet. In diesem Beispiel werden die Größen 128 und 2048 fest codiert. Wenn Sie dies nicht möchten, können Sie die Dimensionen als Parameter an den Aufruf übergeben.

Benutzerdefinierter Anruf auf der GPU

Das benutzerdefinierte GPU-Aufruf-Framework unterscheidet sich etwas von dem auf der CPU. Hier ist ein CUDA-Beispiel, das die gleiche Berechnung von A[i] = B[i % 128] + C[i] wie der obige CPU-Code ausführt.

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

Beachten Sie zunächst, dass die benutzerdefinierte GPU-Aufruffunktion immer noch eine Funktion ist, die auf der CPU ausgeführt wird . Unsere do_custom_call CPU-Funktion ist für das Einreihen der Arbeit an der GPU verantwortlich. Hier wird ein CUDA-Kernel gestartet, aber es kann auch etwas anderes tun, z. B. Cublas aufrufen.

buffers ist ein Array von Zeigern, die sich auf dem Host befinden, und jedes Element, das es enthält, zeigt auf den Gerätespeicher (dh den GPU-Speicher). Die Parameter stehen an erster Stelle, gefolgt vom Ausgabewert. Dies unterscheidet sich erheblich von der CPU-Aufrufkonvention, die zwei Parameter enthält, ins und out . Der Hauptgrund, warum wir auseinander gehen, besteht darin, es möglich zu machen, tupelförmige Ein- / Ausgänge effizient zu handhaben. siehe den folgenden Abschnitt.

Wie im CPU-Beispiel haben wir die Eingabe- und Ausgabepuffergrößen in unserem benutzerdefinierten Aufruf fest codiert. Anders als im CPU-Fall würde es jedoch nicht gut funktionieren, die Puffergrößen als Operanden an den benutzerdefinierten Aufruf zu übergeben. Normalerweise benötigen wir die Puffergrößen, die uns auf der CPU zur Verfügung stehen. Wenn Sie beispielsweise einen Kernel starten, müssen Sie die zu verwendenden Block- / Rasterdimensionen kennen. Wenn wir jedoch die Puffergrößen als Operanden an unseren benutzerdefinierten Aufruf übergeben würden, würden ihre Werte im GPU-Speicher gespeichert. Wir müssten dann zu Beginn unseres Betriebs ein teures synchrones Gerät-zu-Host-Memcpy durchführen, um die Größen zu lesen.

Damit Sie dies umgehen können, stellen wir den opaque Parameter bereit. Sie können dies beim Erstellen des benutzerdefinierten Aufrufs auf eine beliebige Zeichenfolge von Bytes festlegen:

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

Da xla::Shape über eine Protokollpufferdarstellung verfügt, können Sie dieses serialisierte Proto in opaque speichern und in Ihrem benutzerdefinierten GPU-Aufruf deserialisieren. Beachten Sie jedoch, dass sich xla::ShapeProto zwar nicht häufig ändert, sich jedoch ändert. Überprüfen Sie das Git-Protokoll, um festzustellen, wie es sich in der Vergangenheit geändert hat.

Übergeben von Tupeln an benutzerdefinierte Anrufe

Betrachten Sie den folgenden benutzerdefinierten Aufruf.

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

Sowohl auf der CPU als auch auf der GPU wird ein Tupel im Speicher als Array von Zeigern dargestellt. In C ++ - Pseudocode ist der obige Parameter 0 wie folgt aufgebaut.

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

Obwohl die speicherinterne Darstellung von Tupeln in CPU und GPU gleich ist, werden sie in den Konventionen für benutzerdefinierte Aufrufe von CPU und GPU unterschiedlich behandelt.

Tupelausgänge als temporäre Puffer

Tupeleingaben für benutzerdefinierte Anrufe sind praktisch, aber nicht unbedingt erforderlich. Wenn wir keine Tupeleingaben für benutzerdefinierte Aufrufe unterstützen, können Sie die Tupel jederzeit mit get-tuple-element entpacken, bevor Sie sie an den benutzerdefinierten Aufruf übergeben.

Auf der anderen Seite, Tupel Ausgänge lassen Sie Dinge tun , könnten Sie sonst nicht.

Der offensichtliche Grund für Tupelausgaben ist, dass ein benutzerdefinierter Aufruf (oder eine andere XLA-Operation) mehrere unabhängige Arrays zurückgibt.

Weniger offensichtlich ist eine Tupelausgabe auch eine Möglichkeit, Ihren benutzerdefinierten temporären Speicher für Anrufe bereitzustellen. Ja, eine Ausgabe kann einen temporären Puffer darstellen. Beachten Sie, dass ein Ausgabepuffer die Eigenschaft hat, dass die Operation darauf schreiben und nach dem Schreiben daraus lesen kann. Genau das wollen Sie von einem temporären Puffer.

Nehmen wir im obigen Beispiel an, wir wollten den F32[1024] als temporären Puffer verwenden. Dann würden wir das HLO genau wie oben schreiben und niemals den Tupelindex 1 der Ausgabe des benutzerdefinierten Aufrufs lesen.

Tupel in CPU-Custom-Aufrufen

Im CPU-Code haben wir eine Funktion do_custom_call(const void** ins, void* out) . ins ist ein Array mit nur einem Element, das auf param0 . Auf die Unterpuffer von param0 durch Dereferenzieren dieses Zeigers param0 werden, und auf die Unterpuffer von output_tuple durch Dereferenzieren out output_tuple werden.

Tupel in benutzerdefinierten GPU-Aufrufen

Im GPU-Code haben wir eine Funktion do_custom_call(..., void** buffers, ...) . In diesem Fall ist buffers ein Host-Array von sechs Gerätezeigern, einer für jeden Blattpuffer in der Eingabe / Ausgabe. Um die flache Liste zu generieren, durchlaufen wir die Parameter und die Ausgabe und führen für jede eine Vorbestellungsdurchquerung ihrer Form durch. Konkret:

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