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

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

یک تماس سفارشی روی CPU ایجاد کنید

شما می توانید یک دستورالعمل HLO ایجاد کنید که نشان دهنده یک تماس سفارشی از طریق API مشتری XLA است. برای مثال، کد زیر از یک فراخوانی سفارشی برای محاسبه A[i] = B[i % 128]+ C[i] در CPU استفاده می‌کند. (البته شما می توانید – و باید! – این کار را با HLO معمولی انجام دهید.)

#include "xla/client/xla_builder.h"
#include "xla/service/custom_call_target_registry.h"

void do_it() {
  xla::XlaBuilder b("do_it");
  xla::XlaOp param0 =
      xla::Parameter(&b, 0, xla::ShapeUtil::MakeShape(xla::F32, {128}), "p0");
  xla::XlaOp param1 =
      xla::Parameter(&b, 1, xla::ShapeUtil::MakeShape(xla::F32, {2048}), "p1");
  xla::XlaOp custom_call =
      xla::CustomCall(&b, "do_custom_call", /*operands=*/{param0, param1},
                      /*shape=*/xla::ShapeUtil::MakeShape(xla::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 باید ابعاد بافرهایی را که روی آنها کار می کند بداند. در این مثال ما اندازه های 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 = blockIdx.x * blockDim.x + threadIdx.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_t block_dim = 64;
  const int64_t 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 CPU مسئول صف بندی کار روی GPU است. در اینجا یک هسته CUDA راه اندازی می کند، اما می تواند کار دیگری نیز انجام دهد، مانند فراخوانی cuBLAS.

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

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

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

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

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

سیگنال دادن به یک خطا

اگر تماس سفارشی شما با خطا مواجه شد، می‌توانید با استفاده از امضای زیر برای عملکرد خود، خطا را به زمان اجرای XLA (به‌جای خراب شدن یا برگرداندن مطالب مزخرف در بافرهای خروجی) علامت بزنید:

روی CPU:

#include "xla/service/custom_call_status.h"

void do_custom_call(void* out, const void** in, XlaCustomCallStatus* status);

روی پردازنده گرافیکی:

#include "xla/service/custom_call_status.h"

void do_custom_call(CUstream stream, void** buffers, const char* opaque,
                    size_t opaque_len, xla::XlaCustomCallStatus* status);

می‌توانید با استفاده از XlaCustomCallStatusSetFailure علامت خرابی بدهید، به عنوان مثال:

void do_custom_call(void* out, const void** in, XlaCustomCallStatus* status) {
  // ... do some work.

  if (bad_condition) {
    char* error_message = "An error occurred";
    XlaCustomCallStatusSetFailure(status, error_message, strlen(error_message));
    return;
  }

  // ... continue.
}

همچنین می توانید از XlaCustomCallStatusSetSuccess برای نشان دادن موفقیت استفاده کنید، اما XlaCustomCallStatus به طور پیش فرض در وضعیت موفقیت آمیز است، بنابراین نادیده گرفتن کامل آن نیز نشان دهنده موفقیت خواهد بود.

هنگام استفاده از توابع تماس سفارشی با این امضا، باید عملیات custom-call مربوطه را با مجموعه نسخه API مناسب ایجاد کنید، به عنوان مثال:

xla::CustomCall(&b, "do_custom_call", /*operands=*/{param0, param1},
                /*output_shape=*/xla::ShapeUtil::MakeShape(F32, {2048}),
                opaque, /*has_side_effect=*/false,
                /*output_operand_aliasing=*/{}, /*literal=*/nullptr,
                /*schedule=*/xla::CustomCallSchedule::SCHEDULE_NONE,
                /*api_version=*/API_VERSION_STATUS_RETURNING);

در صورت خرابی، هیچ یک از خروجی های تماس سفارشی استفاده نخواهد شد. زمان اجرا XLA محاسبات را خاتمه می دهد. برای محاسبات HLO امکان بازیابی خطا وجود ندارد (مثلاً با گرفتن و مدیریت آن).

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

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

using xla::ShapeUtil;
using xla::F32;
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، یک تاپل در حافظه به عنوان آرایه ای از اشاره گرها نشان داده می شود. در شبه کد 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;

اگرچه نمایش داخل حافظه تاپل ها در CPU و GPU یکسان است، اما در قراردادهای فراخوانی سفارشی CPU و GPU با آنها متفاوت رفتار می شود.

خروجی ها را به عنوان بافرهای دمایی چند برابر کنید

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

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

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

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

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

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

در کد CPU، یک تابع do_custom_call(const void** ins, void* out) داریم. ins آرایه ای با تنها یک عنصر است که به param0 اشاره می کند. بافرهای فرعی param0 با عدم ارجاع آن اشاره گر و بافرهای فرعی 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