ترجمت واجهة Cloud Translation API‏ هذه الصفحة.
Switch to English

مكالمات XLA المخصصة

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

مكالمة مخصصة على وحدة المعالجة المركزية

يمكنك إنشاء تعليمات HLO التي تمثل مكالمة مخصصة عبر واجهة برمجة تطبيقات عميل 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 تحتاج إلى معرفة أبعاد المخازن المؤقتة التي تعمل عليها. في هذا المثال ، نقوم بعمل كود ثابت للأحجام 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 لا تزال وظيفة يتم تنفيذها على وحدة المعالجة المركزية . إن وظيفة do_custom_call CPU مسؤولة عن متابعة العمل على GPU. هنا تطلق نواة CUDA ، ولكن يمكنها أيضًا القيام بشيء آخر ، مثل استدعاء cublas.

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

كما هو الحال في مثال وحدة المعالجة المركزية ، قمنا بترميز أحجام المخزن المؤقت للإدخال والإخراج في مكالمتنا المخصصة. ولكن على عكس حالة وحدة المعالجة المركزية ، فإن تمرير أحجام المخزن المؤقت كمعاملين إلى المكالمة المخصصة لن يعمل بشكل جيد. عادة ما نحتاج إلى أحجام المخزن المؤقت المتاحة لنا على وحدة المعالجة المركزية ؛ على سبيل المثال عند إطلاق نواة ، نحتاج إلى معرفة أبعاد الكتلة / الشبكة لاستخدامها. ولكن إذا تمكنا من تمرير أحجام المخزن المؤقت مثل المعاملات إلى مكالمتنا المخصصة ، فستعيش قيمها في ذاكرة 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 وإلغاء تسلسله داخل المكالمة المخصصة لوحدة معالجة الرسومات. لاحظ أنه على الرغم من أن xla::ShapeProto لا يتغير بشكل متكرر ، إلا أنه يتغير. تحقق من سجل git لمعرفة كيف تغير في الماضي.

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

خذ بعين الاعتبار المكالمة المخصصة التالية.

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

في كل من وحدة المعالجة المركزية ووحدة معالجة الرسومات ، يتم تمثيل مجموعة في الذاكرة كمجموعة من المؤشرات. في C ++ - pseudocode ، تم وضع المعلمة 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;
 

على الرغم من أن تمثيل tuple في الذاكرة هو نفسه في وحدة المعالجة المركزية ووحدة معالجة الرسومات ، إلا أنه يتم التعامل معها بشكل مختلف في اصطلاحات المكالمات المخصصة لوحدة المعالجة المركزية ووحدة معالجة الرسومات.

نواتج Tuple كمخازن مؤقتة

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

من ناحية أخرى ، تتيح لك مخرجات الصفوف القيام بأشياء لا يمكنك خلاف ذلك.

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

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

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

Tuples في مكالمات مخصصة CPU

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

Tuples في مكالمات 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