Bu sayfa, Cloud Translation API ile çevrilmiştir.
Switch to English

XLA Özel Çağrılar

Bu belgede, XLA "özel çağrılarının" nasıl yazılacağı ve kullanılacağı açıklanmaktadır. Özel çağrılar, bir XLA programından C ++ veya CUDA gibi bir programlama dilinde yazılmış kodu çağırmanızı sağlar.

CPU'da özel çağrı

XLA'nın istemci API'sı aracılığıyla özel bir çağrıyı temsil eden bir HLO talimatı oluşturabilirsiniz. Bu, yazılı olarak TensorFlow aracılığıyla gösterilmez.

Örneğin, aşağıdaki kod CPU'da A[i] = B[i % 128] + C[i] hesaplamak için özel bir çağrı kullanır. (Elbette bunu düzenli HLO ile yapabilirsin ve yapmalısın!)

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

do_custom_call fonksiyonunun üzerinde çalıştığı tamponların boyutlarını do_custom_call gerektiğine dikkat edin. Bu örnekte 128 ve 2048 boyutlarını sabit kodluyoruz. Bunu yapmak istemiyorsanız, boyutları aramaya parametre olarak iletebilirsiniz.

GPU'da özel çağrı

GPU özel çağrı çerçevesi CPU'dekinden biraz farklıdır. Yukarıdaki CPU kodu ile aynı A[i] = B[i % 128] + C[i] hesaplamasını yapan bir CUDA örneği.

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

İlk olarak GPU özel çağrı işlevinin hala CPU'da yürütülen bir işlev olduğuna dikkat edin. do_custom_call CPU fonksiyonumuz GPU üzerinde çalışmayı sıralamaktan sorumludur. Burada bir CUDA çekirdeği piyasaya sürüyor, ancak çağrı küpleri gibi başka bir şey de yapabilir.

buffers , ana bilgisayarda yaşayan bir işaretçiler dizisidir ve içerdiği her öğe, aygıt (yani GPU) belleğine puan verir. Parametreler önce gelir, ardından çıktı değeri gelir. Bu iki parametreler, vardır CPU çağrı kuralı, oldukça farklıdır ins ve out . Ayırmamızın ana nedeni, tuple şeklindeki girişlerin / çıkışların verimli bir şekilde kullanılmasını mümkün kılmaktır; aşağıdaki bölüme bakın.

CPU örneğinde olduğu gibi, giriş ve çıkış arabellek boyutlarını özel aramamıza sabit olarak kodladık. Ancak CPU durumundan farklı olarak, arabellek boyutlarını işlenen olarak özel aramaya geçirmek iyi çalışmaz. Genellikle CPU'da kullanabileceğimiz tampon boyutlarına ihtiyacımız var; örneğin bir çekirdeği başlatırken, kullanılacak blok / ızgara boyutlarını bilmemiz gerekir. Ancak arabellek boyutlarını işlenen olarak özel çağrımıza geçirecek olsaydık, değerleri GPU belleğinde yaşayacaktı. Daha sonra, sadece boyutları okumak için operasyonumuzun başında pahalı bir senkronize cihazdan ana bilgisayara memcpy yapmamız gerekir.

Bu sorunu çözmenize izin vermek için opaque parametreyi sağlıyoruz. Özel aramayı oluştururken bunu keyfi bir bayt dizesine ayarlayabilirsiniz:

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

xla::Shape bir protokol tamponu temsiline sahip olduğundan, bu serileştirilmiş protokolü opaque içinde saklayabilir ve GPU özel çağrınızda serisini kaldırabilirsiniz. Ancak xla::ShapeProto sık sık değişmese de değiştiğine dikkat edin. Geçmişte nasıl değiştiğini görmek için git günlüğünü kontrol edin.

Özel çağrılara tuples aktarma

Aşağıdaki özel aramayı düşünün.

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

Hem CPU hem de GPU'da, bir demet bellekte bir işaretçi dizisi olarak temsil edilir. C ++ - sözde kodunda, yukarıdaki parametre 0 aşağıdaki gibi düzenlenir.

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

Her ne kadar tupleslerin bellek içi gösterimi CPU ve GPU'da aynı olsa da, CPU ve GPU özel çağrı arama kurallarında farklı şekilde ele alınır.

Geçici tamponlar olarak grup çıkışları

Özel çağrılara grup girişleri kolaylık sağlar, ancak kesinlikle gerekli değildir. Özel aramalarda grup girişlerini desteklemiyorsak, özel aramaya geçmeden önce tuples'ı get-tuple-element kullanarak açabilirsiniz.

Öte yandan, grup çıktıları başka türlü yapamayacağınız şeyleri yapmanızı sağlar.

Grup çıkışlarına sahip olmanın bariz nedeni, özel bir aramanın (veya başka bir XLA opunun) çoklu bağımsız dizileri bu şekilde döndürmesidir.

Ancak daha az belli ki, bir grup çıkışı aynı zamanda özel çağrı temp belleğinizi vermenin bir yoludur. Evet, bir çıktı geçici arabelleği temsil edebilir. Bir çıkış arabelleğinin, op'un ona yazabileceği özelliğe sahip olduğunu ve yazıldıktan sonra ondan okuyabileceğini düşünün. Geçici bir arabellekten tam olarak istediğiniz budur.

Yukarıdaki örnekte, F32[1024] i geçici bir tampon olarak kullanmak istediğimizi varsayalım. Sonra HLO'yu yukarıdaki gibi yazardık ve özel çağrının çıktısının 1 numaralı dizinini asla okumazdık.

CPU özel çağrılarında tuples

CPU kodunda, do_custom_call(const void** ins, void* out) . ins , param0 işaret eden tek bir öğeye sahip bir dizidir. Arasında subbuffers param0 bu işaretçiyi kaldırma tarafından erişilebilir ve subbuffers output_tuple kaldırma tarafından erişilebilir out .

GPU özel çağrılarındaki tuples

GPU kodunda, do_custom_call(..., void** buffers, ...) . Bu durumda buffers , giriş / çıkıştaki her bir yaprak arabelleği için bir tane olmak üzere altı aygıt işaretçisi içeren bir ana bilgisayar dizisidir. Düz listeyi oluşturmak için, parametreler ve çıktılar üzerinde tekrarlar ve her biri için şeklinin ön geçişini yaparız. somut olarak:

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