المكالمات المخصّصة عبر XLA

يصف هذا المستند كيفية كتابة واستخدام استدعاءات XLA المخصصة. تتيح لك الاستدعاءات المخصّصة استدعاء تعليمة برمجية مكتوبة بلغة برمجة مثل C++ أو CUDA من برنامج XLA.

إنشاء مكالمة مخصّصة على وحدة المعالجة المركزية (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}));
}

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. إذا كنت لا تريد إجراء ذلك، يمكنك ضبط السمات كمعلَمات للمكالمة.

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

يختلف إطار عمل الاتصال المخصّص لوحدة معالجة الرسومات إلى حد ما عن إطار عمل الاتصال المخصّص في وحدة المعالجة المركزية (CPU). وفي ما يلي مثال على 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, 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");

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

buffers هي مصفوفة من المؤشرات المتوفّرة على المضيف، ويحتوي كل عنصر فيه على نقاط إلى ذاكرة الجهاز (أي وحدة معالجة الرسومات). تأتي المعلمات أولاً، تليها قيمة المخرجات. ويختلف هذا بشكل ملحوظ عن اصطلاح استدعاء وحدة المعالجة المركزية (CPU)، الذي يحتوي على مَعلمتَين، ins وout. ويجعل اصطلاح الاتصال عبر وحدة معالجة الرسومات من الممكن التعامل مع الإدخالات/المخرجات على شكل صف بكفاءة.

وكما في مثال وحدة المعالجة المركزية (CPU)، قمنا بترميز أحجام المخزن المؤقت للإدخال والإخراج في مكالمتنا المخصصة. وعلى عكس حالة وحدة المعالجة المركزية (CPU)، فإن تمرير أحجام المخزن المؤقت كمعاملات للاستدعاء المخصص لن يعمل بشكل جيد. نحتاج عادةً إلى أحجام الموارد الاحتياطية المتاحة لنا على وحدة المعالجة المركزية (CPU) (على سبيل المثال، عند تشغيل نواة نواة، نحتاج إلى معرفة أبعاد الكتلة/الشبكة المراد استخدامها). ولكن إذا قمنا بتمرير أحجام المخزن المؤقت كمعاملات للاستدعاء المخصص، فستبقى قيمها في ذاكرة وحدة معالجة الرسومات. في بداية العملية، علينا إجراء عملية 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 وإلغاء تسلسله في طلب وحدة معالجة الرسومات المخصّص. يُرجى العلم أنّه على الرغم من أنّ السمة 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 المقابلة باستخدام مجموعة إصدارات واجهة برمجة التطبيقات المناسبة، مثلاً:

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

إخراج الصف كوحدات تخزين مؤقتة

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

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

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

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

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

صفوف في المكالمات المخصصة بوحدة المعالجة المركزية (CPU)

في رمز وحدة المعالجة المركزية، لدينا الدالة do_custom_call(const void** ins, void* out). ins هي صفيف يحتوي على عنصر واحد فقط ويشير إلى param0. يمكن الوصول إلى الموارد الاحتياطية لـ param0 من خلال إلغاء الإشارة إلى هذا المؤشر، ويمكن الوصول إلى الموارد الاحتياطية لـ output_tuple من خلال إلغاء الإشارة إلى out.

صفوف في المكالمات المخصصة لوحدة معالجة الرسومات

في رمز وحدة معالجة الرسومات، لدينا الدالة 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