تاریخ را ذخیره کنید! Google I / O 18-20 مه بازمی گردد اکنون ثبت نام کنید
این صفحه به‌وسیله ‏Cloud Translation API‏ ترجمه شده است.
Switch to English

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

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

تماس سفارشی با پردازنده مرکزی

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

به عنوان مثال ، کد زیر از یک تماس سفارشی برای محاسبه A[i] = B[i % 128] + C[i] روی پردازنده استفاده می کند. (البته شما می توانید - و باید! - این کار را با 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 نسبت به پردازنده تا حدودی متفاوت است. در اینجا یک مثال 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 اجرا می شود . عملکرد پردازنده do_custom_call ما وظیفه ایجاد کار روی GPU را بر عهده دارد. در اینجا هسته CUDA را راه اندازی می کند ، اما همچنین می تواند کار دیگری مانند تماس با کوباها را انجام دهد.

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

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

برای اینکه بتوانید در این مورد کار کنید ، پارامتر 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::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 به طور متفاوتی رفتار می شود.

خروجی های Tuple به عنوان بافرهای temp

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

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

دلیل واضح داشتن خروجی های tuple این است که چگونه یک تماس سفارشی (یا هر عمل دیگر XLA) چندین آرایه مستقل را برمی گرداند.

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

در مثال بالا ، فرض کنید می خواهیم از F32[1024] به عنوان بافر temp استفاده کنیم. سپس ما می توانیم HLO را همانند بالا بنویسیم و هرگز نمیتوانیم شاخص tuple 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