Halaman ini diterjemahkan oleh Cloud Translation API.
Switch to English

Panggilan XLA Kustom

Dokumen ini menjelaskan bagaimana menulis dan menggunakan XLA "kustom panggilan". panggilan khusus memungkinkan Anda menjalankan kode yang ditulis dalam bahasa pemrograman seperti C ++ atau CUDA dari program XLA.

Custom-panggilan pada CPU

Anda dapat membuat instruksi HLO yang merupakan custom-panggilan melalui XLA ini klien API. Ini tidak terkena via TensorFlow sebagai tulisan.

Sebagai contoh, kode berikut menggunakan custom-panggilan untuk menghitung A[i] = B[i % 128] + C[i] pada CPU. (Tentu saja Anda bisa - dan harus - melakukan hal ini dengan HLO biasa.)

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

Perhatikan bahwa fungsi do_custom_call perlu mengetahui dimensi dari buffer beroperasi lebih. Dalam contoh ini kita hardcode ukuran 128 dan 2048. Jika Anda tidak ingin melakukan ini, Anda dapat melewati dimensi sebagai parameter untuk panggilan.

Custom-panggilan pada GPU

Kerangka panggilan GPU kustom agak berbeda dari yang pada CPU. Berikut ini adalah contoh CUDA yang melakukan hal yang sama A[i] = B[i % 128] + C[i] perhitungan sebagai kode CPU di atas.

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

Perhatikan pertama bahwa GPU fungsi panggilan kustom masih fungsi dijalankan pada CPU. Kami do_custom_call fungsi CPU bertanggung jawab untuk enqueueing bekerja pada GPU. Di sini meluncurkan kernel CUDA, tetapi juga bisa melakukan sesuatu yang lain, seperti panggilan cublas.

buffers adalah array dari pointer yang hidup pada host, dan setiap elemen mengandung poin ke perangkat (yaitu GPU) memori. Parameter datang pertama, diikuti oleh nilai output. Hal ini terutama berbeda dari CPU konvensi pemanggilan, yang memiliki dua params, ins dan out . Alasan utama kita menyimpang adalah untuk memungkinkan untuk menangani tupel berbentuk input / output secara efisien; lihat bagian di bawah.

Seperti pada contoh CPU, kita sudah hardcoded input dan output buffer ukuran dalam panggilan kebiasaan kami. Namun tidak seperti dalam kasus CPU, melewati buffer ukuran sebagai operan ke panggilan kustom tidak akan bekerja dengan baik. Biasanya kita perlu buffer ukuran yang tersedia kepada kami pada CPU; misalnya ketika meluncurkan sebuah kernel, kita perlu mengetahui dimensi blok / grid untuk digunakan. Tapi jika kita lulus ukuran buffer sebagai operan panggilan kebiasaan kami, nilai-nilai mereka akan tinggal di memori GPU. Kami kemudian harus melakukan sinkron memcpy perangkat-to-host mahal pada awal operasi kami hanya untuk membaca ukuran.

Untuk membiarkan Anda bekerja di sekitar ini, kami menyediakan opaque parameter. Anda dapat mengatur ini untuk string sewenang-wenang byte ketika Anda membuat panggilan kustom:

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

Sejak xla::Shape memiliki representasi protokol penyangga, Anda bisa menyimpan ini di dalam proto serial dari opaque dan deserialize dalam GPU custom-panggilan Anda. Namun perlu dicatat bahwa meskipun xla::ShapeProto tidak sering berubah, itu tidak berubah. Periksa log git untuk melihat bagaimana ia telah berubah di masa lalu.

Melewati tupel untuk custom-panggilan

Pertimbangkan hal berikut custom-panggilan.

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

Pada kedua CPU dan GPU, tupel direpresentasikan dalam memori sebagai array dari pointer. Dalam C ++ - pseudocode, parameter 0 atas ditata sebagai berikut.

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

Meskipun di memori representasi dari tupel adalah sama dalam CPU dan GPU, mereka ditangani secara berbeda dalam CPU dan GPU custom-panggilan menelepon konvensi.

Tupel output sebagai buffer suhu

masukan tupel untuk custom-panggilan kenyamanan, tetapi mereka tidak benar-benar diperlukan. Jika kita tidak mendukung input tuple panggilan kustom, Anda selalu bisa membongkar tupel menggunakan get-tuple-elemen sebelum melewati mereka untuk panggilan kustom.

Di sisi lain, tuple output membiarkan Anda melakukan hal-hal yang Anda tidak bisa sebaliknya.

Alasan yang jelas untuk memiliki output tupel adalah, itu bagaimana panggilan kustom (atau op XLA lainnya) mengembalikan beberapa array independen.

Tapi kurang jelas, output tuple juga merupakan cara untuk memberikan kustom memori panggilan temp Anda. Ya, output dapat mewakili buffer temp. Pertimbangkan, sebuah output buffer memiliki properti yang op dapat menulis untuk itu, dan dapat membaca dari itu setelah sudah ditulis. Itulah apa yang Anda inginkan dari buffer temp.

Dalam contoh di atas, misalkan kita ingin menggunakan F32[1024] sebagai penyangga temp. Kemudian kita akan menulis HLO hanya seperti di atas, dan kami akan hanya pernah membaca tuple indeks 1 output panggilan kustom ini.

Tupel dalam CPU custom-panggilan

Dalam kode CPU, kita memiliki fungsi do_custom_call(const void** ins, void* out) . ins adalah array dengan hanya satu elemen, yang menunjuk ke param0 . The subbuffers dari param0 dapat diakses oleh dereferencing pointer itu, dan subbuffers dari output_tuple dapat diakses oleh dereferencing out .

Tupel di GPU custom-panggilan

Dalam kode GPU, kita memiliki fungsi do_custom_call(..., void** buffers, ...) . Dalam hal ini buffers adalah array tuan rumah enam pointer perangkat, satu untuk setiap buffer daun di input / output. Untuk menghasilkan daftar datar, kita iterate atas parameter dan output, dan untuk setiap kita melakukan preorder traversal dari bentuknya. Konkretnya:

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