این صفحه به‌وسیله ‏Cloud Translation API‏ ترجمه شده است.
Switch to English

تماسهای سفارشی XLA

این سند نحوه نوشتن و استفاده از "تماسهای سفارشی" XLA را شرح می دهد. تماسهای سفارشی به شما امکان می دهد از برنامه XLA کدی که به زبان برنامه نویسی مانند C ++ یا CUDA نوشته شده استناد کنید.

سفارشی با CPU تماس بگیرید

می توانید یک دستورالعمل HLO ایجاد کنید که از طریق API مشتری XLA یک تماس سفارشی را نشان می دهد. این امر از طریق نوشتن TensorFlow در معرض دید قرار نمی گیرد.

به عنوان مثال ، کد زیر از یک تماس اختصاصی برای محاسبه A[i] = B[i % 128] + C[i] در CPU استفاده می کند. (البته شما می توانید - و باید! - این کار را با 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");
 

توجه کنید که عملکرد do_custom_call باید ابعاد do_custom_call را که در آن کار می کند ، بداند. در این مثال اندازه های 128 و 2048 را به صورت کدکد می زنیم. اگر نمی خواهید این کار را انجام دهید ، می توانید ابعاد را به عنوان پارامترها به تماس منتقل کنید.

تماس سفارشی در GPU

چارچوب تماس های سفارشی GPU تا حدودی متفاوت از پردازنده CPU است. در اینجا یک مثال CUDA وجود دارد که محاسبه A[i] = B[i % 128] + C[i] به عنوان کد CPU فوق انجام می دهد.

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

ابتدا توجه کنید که عملکرد تماس GPU سفارشی هنوز تابعی است که بر روی CPU اجرا می شود . عملکرد CPU do_custom_call ما وظیفه تنظیم کارهای مربوط به پردازنده گرافیکی را دارد. در اینجا هسته CUDA را راه اندازی می کند ، اما می تواند کار دیگری انجام دهد ، مانند مکالمه های فراخوانی.

buffers مجموعه ای از نشانگرهاست که روی میزبان زندگی می کند و هر عنصر در آن دارای نقاطی به حافظه (یعنی GPU) حافظه است. پارامترها ابتدا آمده و به دنبال آن مقدار مقدار خروجی است. این امر به ویژه با کنوانسیون فراخوانی CPU که دارای دو param ، ins و out . دلیل اصلی که ما واگرایی می کنیم این است که امکان کارآمد ورود و خروجی های توپی شکل را به صورت کارآمد فراهم می کنیم. بخش زیر را ببینید.

همانطور که در مثال CPU ، اندازه های بافر ورودی و خروجی را در تماس سفارشی خود سخت کده ایم. اما برخلاف مورد CPU ، انتقال اندازه بافر به عنوان اپراتورها به تماس سفارشی خوب نخواهد بود. معمولاً به اندازه های بافر موجود در CPU نیاز داریم. به عنوان مثال هنگام راه اندازی یک هسته ، ما باید از ابعاد بلاک / شبکه برای استفاده بدانیم. اما اگر بخواهیم اندازه های بافر را به عنوان operands به تماس اختصاصی خود منتقل کنیم ، مقادیر آنها در حافظه GPU زندگی می کنند. در ادامه می خواهیم فقط برای خواندن اندازه ها ، یک ممکپی دستگاه همزمان و میزبان گرانقیم را برای شروع کار خود انجام دهیم.

برای اینکه بتوانید در این زمینه کار کنید ، پارامتر opaque را ارائه می دهیم. می توانید هنگام برقراری تماس سفارشی ، این کار را روی یک بایت دلخواه بایت تنظیم کنید:

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

از آنجا که xla::Shape دارای یک پروتئین بافر پروتکل است ، می توانید این نسخه سریال شده را در داخل opaque و آن را در تماس گرافیکی GPU خود به xla::Shape کنید. توجه داشته باشید که اگرچه xla::ShapeProto به طور مکرر تغییر نمی کند ، اما تغییر می کند. ورود به سیستم git را بررسی کنید تا ببینید چگونه در گذشته تغییر کرده است.

انتقال تاپلها به تماسهای سفارشی

تماس سفارشی زیر را در نظر بگیرید.

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

در هر دو CPU و GPU ، یک tuple به عنوان آرایه ای از نشانگرها در حافظه نمایش داده می شود. در C ++ - شبه کد ، پارامتر 0 در بالا به شرح زیر است.

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

اگرچه نمایش حافظه از tuples در CPU و GPU یکسان است ، اما در کنوانسیون های تماس تلفنی با CPU و GPU به طور متفاوتی به کار گرفته می شوند.

خروجی ها را به عنوان بافر فشار دهید

ورودی های ورودی به تماس های سفارشی یک راحتی هستند ، اما آنها کاملاً ضروری نیستند. اگر ما از ورودی های تاپل در تماس های سفارشی پشتیبانی نمی کردیم ، همیشه می توانید قبل از انتقال آنها به تماس سفارشی ، از عناصر get-tuple استفاده کنید.

از طرف دیگر ، خروجی های تاپل به شما امکان می دهند کارهایی را انجام دهید که در غیر این صورت نمی توانستید.

دلیل بارز خروجی های تاپل این است که تماس تلفنی (یا هر XLA دیگر) چندین آرایه مستقل را برمی گرداند.

اما کمتر واضح است ، یک خروجی tuple نیز راهی برای دادن حافظه دماغ تماس سفارشی شماست. بله ، یک خروجی می تواند یک بافر دما باشد. در نظر بگیرید ، یک بافر خروجی خاصیتی دارد که op می تواند برای آن بنویسد ، و می تواند پس از نوشتن آن ، آن را بخواند. این دقیقاً همان چیزی است که شما از بافر دما می خواهید.

در مثال بالا ، فرض کنید می خواهیم از F32[1024] به عنوان بافر دما استفاده کنیم. سپس ما می خواهیم HLO را دقیقاً مانند بالا بنویسیم ، و هیچ وقت نمی خواستیم شاخص تاپل 1 خروجی تماس سفارشی را بخوانیم.

تماسهای سفارشی CPU را لمس کنید

در کد CPU ، یک تابع do_custom_call(const void** ins, void* out) . ins یک آرایه با تنها یک عنصر است که به param0 اشاره می param0 . subbuffers از param0 بوسیله ارجاع که اشاره گر در دسترس هستند، و subbuffers از output_tuple بوسیله ارجاع در دسترس هستند out .

تماس های گرافیکی GPU را لمس کنید

در کد GPU ، ما یک تابع do_custom_call(..., void** buffers, ...) . در این حالت buffers یک میزبان از شش نشانگر دستگاه است ، یکی برای هر بافر برگ در ورودی / خروجی. برای تولید لیست صاف ، ما بیش از پارامترها و خروجی تکرار می کنیم ، و برای هر یک یک پیشروی پیش شکل از شکل آن را انجام می دهیم. بطور مشخص:

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