مكالمات مخصّصة بتنسيق XLA

يوضّح هذا المستند طريقة كتابة واستخدام طلبات XLA المخصّصة باستخدام مكتبة XLA FFI. الاستدعاء المخصّص هو آلية لوصف "عملية" خارجية في وحدة HLO إلى المحول البرمجي XLA (في وقت التجميع)، في حين أنّ XLA FFI هي آلية لتسجيل تنفيذ هذه العمليات باستخدام XLA (في وقت التشغيل). يعني الاختصار FFI "واجهة الدوال الخارجية"، وهو مجموعة من واجهات برمجة التطبيقات C التي تحدد واجهة ثنائية (ABI) لاستدعاء XLA إلى رمز خارجي مكتوب بلغات برمجة أخرى. توفّر XLA روابط رأس فقط لـ XLA FFI مكتوبة بلغة C++ ، ما يخفي جميع التفاصيل المنخفضة المستوى لواجهات برمجة تطبيقات C الأساسية عن المستخدم النهائي.

إنشاء مكالمة مخصّصة في وحدة المعالجة المركزية (CPU)

يمكنك إنشاء تعليمات HLO التي تمثل مكالمة مخصصة عبر واجهة برمجة تطبيقات عميل 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}),
        /*opaque=*/"", /*has_side_effect=*/false,
        /*output_operand_aliasing=*/{}, /*literal=*/nullptr,
        /*schedule=*/CustomCallSchedule::SCHEDULE_NONE,
        /*api_version=*/CustomCallApiVersion::API_VERSION_TYPED_FFI);
}

// Constrain custom call arguments to rank-1 buffers of F32 data type.
using BufferF32 = xla::ffi::BufferR1<xla::ffi::DataType::F32>;

// Implement a custom call as a C+ function. Note that we can use `Buffer` type
// defined by XLA FFI that gives us access to buffer data type and shape.
xla::ffi::Error do_custom_call(BufferF32 in0, BufferF32 in1,
                               xla::ffi::Result<BufferF32> out) {
  size_t d0 = in0.dimensions[0];
  size_t d1 = in1.dimensions[0];

  // Check that dimensions are compatible.
  assert(out->dimensions[0] == d1 && "unexpected dimensions");

  for (size_t i = 0; i < d1; ++i) {
    out->data[i] = in0.data[i % d0] + in1.data[i];
  }
}

// Explicitly define an XLA FFI handler signature and bind it to the
// `do_custom_call` implementation. XLA FFI handler can automatically infer
// type signature from the custom call function, but it relies on magical
// template metaprogramming an explicit binding provides and extra level of
// type checking and clearly states custom call author intentions.
XLA_FFI_DEFINE_HANDLER(handler, do_custom_call,
                       ffi::Ffi::Bind()
                           .Arg<Buffer>()
                           .Arg<Buffer>()
                           .Ret<Buffer>());

// Registers `handler` with and XLA FFI on a "Host" platform.
XLA_FFI_REGISTER_HANDLER(xla::ffi::GetXlaFfiApi(), "do_custom_call",
                         "Host", handler);

إنشاء مكالمة مخصّصة في وحدة معالجة الرسومات

تسجيل المكالمات المخصصة لوحدة معالجة الرسومات مع XLA FFI متطابق تقريبًا، والفرق الوحيد هو أنه بالنسبة إلى وحدة معالجة الرسومات التي تحتاج إلى طلب بث أساسي للنظام الأساسي (بث CUDA أو ROCM) لتتمكن من تشغيل النواة على الجهاز. إليك مثال على CUDA يقوم بنفس طريقة الحوسبة (A[i] = B[i % 128] + C[i]) التي يقوم بها رمز وحدة المعالجة المركزية أعلاه.

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, BufferF32 in0, BufferF32 in1,
                    xla::ffi::Result<BufferF32> out) {
  size_t d0 = in0.dimensions[0];
  size_t d1 = in1.dimensions[0];
  size_t d2 = out->dimensions[0];

  assert(d0 == 128 && d1 == 2048 && d2 == 2048 && "unexpected dimensions");

  const int64_t block_dim = 64;
  const int64_t grid_dim = 2048 / block_dim;
  custom_call_kernel<<<grid_dim, block_dim, 0, stream>>>(
    in0.data, in1.data, out->data);
}

XLA_FFI_DEFINE_HANDLER(handler, do_custom_call,
                       ffi::Ffi::Bind()
                           .Ctx<xla::ffi::PlatformStream<CUstream>>()
                           .Arg<BufferF32>()
                           .Arg<BufferF32>()
                           .Ret<BufferF32>());

XLA_FFI_REGISTER_HANDLER(xla::ffi::GetXlaFfiApi(), "do_custom_call",
                         "CUDA", handler);

تجدر الإشارة أولاً إلى أنّ وظيفة الطلب المخصّصة لوحدة معالجة الرسومات لا تزال وظيفة يتم تنفيذها على وحدة المعالجة المركزية (CPU). تكون وظيفة وحدة المعالجة المركزية do_custom_call مسؤولة عن إدراج العمل في قائمة الانتظار على وحدة معالجة الرسومات. هنا تُطلق نواة CUDA، لكن يمكنها أيضًا تنفيذ شيء آخر، مثل استدعاء دالة cuBLAS.

توجد الوسيطات والنتائج أيضًا على المضيف، ويحتوي عضو البيانات على مؤشر إلى ذاكرة جهاز (أي وحدة معالجة الرسومات). يكون للمخزن المؤقت الذي يتم تمريره إلى معالج الاستدعاء المخصص شكل المخازن الاحتياطية الأساسية للجهاز، بحيث يمكن للاستدعاء المخصص حساب معلمات تشغيل النواة منها.

تمرير الصفوف إلى المكالمات المخصصة

ضع في الاعتبار الاستدعاء المخصص التالي.

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) ووحدة معالجة الرسومات، يتم تمثيل الصف في الذاكرة كمجموعة من المؤشرات. عندما يستدعي XLA استدعاءات مخصصة باستخدام وسيطات صف أو نتائج، يتم تبسيطها وتمريرها كوسيطات أو نتائج للمخزن المؤقت العادي.

مخرجات الصفوف كموارد مؤقتة

تعتبر المدخلات الصفية للمكالمات المخصصة أمرًا سهلاً، ولكنها ليست ضرورية للغاية. إذا لم نكن ندعم مدخلات الصف للاستدعاءات المخصصة، فيمكنك دائمًا فك ضغط الصفوف باستخدام عنصر الحصول على الصفوف قبل تمريرها إلى الاستدعاء المخصص.

من ناحية أخرى، تتيح لك المخرجات في الصف تنفيذ مهام لا تستطيع بخلافها.

السبب الواضح للحصول على مخرجات الصف هو أن مخرجات الصف هي كيفية عرض استدعاء مخصص (أو أي عملية XLA أخرى) لصفائف مستقلة متعددة.

ولكن أقل وضوحًا، فإن ناتج الصف يعد أيضًا وسيلة لمنح ذاكرة مؤقتة للمكالمة المخصصة. نعم، يمكن أن يمثل الناتج موردًا مؤقتًا. ضع في اعتبارك أن المخزن المؤقت للمخرجات لديه الخاصية التي يمكن للعمل الكتابة عليها، ويمكن القراءة منه بعد الكتابة عليها. هذا بالضبط ما تريده من التخزين المؤقت.

في المثال أعلاه، لنفترض أنّنا أردنا استخدام F32[1024] كمخزن مؤقّت. ثم نكتب HLO كما هو موضح أعلاه، ولن نقرأ مطلقًا فهرس الصف 1 من مخرجات الاستدعاء المخصص.