এই ডকুমেন্টটি বর্ণনা করে কিভাবে XLA কাস্টম কল লিখতে এবং ব্যবহার করতে হয়। কাস্টম কলগুলি আপনাকে একটি XLA প্রোগ্রাম থেকে C++ বা CUDA এর মতো একটি প্রোগ্রামিং ভাষায় লিখিত কোড আহ্বান করতে দেয়।
CPU-তে একটি কাস্টম কল তৈরি করুন
আপনি একটি HLO নির্দেশ তৈরি করতে পারেন যা XLA এর ক্লায়েন্ট API এর মাধ্যমে একটি কাস্টম কল উপস্থাপন করে। উদাহরণস্বরূপ, CPU-তে A[i] = B[i % 128]+ C[i] গণনা করতে নিম্নলিখিত কোডটি একটি কাস্টম কল ব্যবহার করে। (অবশ্যই আপনি করতে পারেন - এবং করা উচিত! - নিয়মিত 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 উদাহরণ রয়েছে যা উপরের CPU কোডের মতো একই গণনা ( 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");
প্রথমে লক্ষ্য করুন যে জিপিইউ কাস্টম কল ফাংশনটি এখনও সিপিইউতে কার্যকর করা একটি ফাংশন । do_custom_call CPU ফাংশন GPU-তে কাজ সারিবদ্ধ করার জন্য দায়ী। এখানে এটি একটি CUDA কার্নেল চালু করে, তবে এটি অন্য কিছুও করতে পারে, যেমন cuBLAS কল।
buffers হল পয়েন্টারগুলির একটি অ্যারে যা হোস্টে থাকে এবং প্রতিটি উপাদান এতে ডিভাইসের (যেমন GPU) মেমরির পয়েন্ট থাকে। পরামিতিগুলি প্রথমে আসে, তারপরে আউটপুট মান। এটি উল্লেখযোগ্যভাবে CPU কলিং কনভেনশন থেকে আলাদা, যার দুটি প্যারাম আছে, ins এবং out । GPU কলিং কনভেনশন টিপল-আকৃতির ইনপুট/আউটপুট দক্ষতার সাথে পরিচালনা করা সম্ভব করে তোলে।
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 ঘন ঘন পরিবর্তন হয় না, তবে এটি পরিবর্তিত হয় । অতীতে এটি কীভাবে পরিবর্তিত হয়েছে তা দেখতে গিট লগটি দেখুন।
একটি ত্রুটি সংকেত
যদি আপনার কাস্টম কল একটি ত্রুটির সম্মুখীন হয়, তাহলে আপনি আপনার ফাংশনের জন্য নিম্নলিখিত স্বাক্ষর ব্যবহার করে XLA রানটাইমে (যেমন আউটপুট বাফারে ক্র্যাশ বা ননসেন্স রিটার্ন করার পরিবর্তে) ত্রুটিটি সংকেত দিতে পারেন:
CPU-তে:
#include "xla/service/custom_call_status.h"
void do_custom_call(void* out, const void** in, XlaCustomCallStatus* status);
GPU-তে:
#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 ডিফল্টভাবে একটি সফল অবস্থায় রয়েছে, তাই এটিকে সম্পূর্ণরূপে উপেক্ষা করাও সাফল্যের ইঙ্গিত দেবে।
এই স্বাক্ষরের সাথে কাস্টম কল ফাংশন ব্যবহার করার সময়, আপনাকে অবশ্যই উপযুক্ত API সংস্করণ সেটের সাথে সংশ্লিষ্ট 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 রানটাইম গণনা বন্ধ করবে। একটি এইচএলও কম্পিউটেশনের পক্ষে ত্রুটি থেকে পুনরুদ্ধার করা সম্ভব নয় (যেমন এটি ধরা এবং পরিচালনা করে)।
কাস্টম কলে টিপল পাস করা
নিম্নলিখিত কাস্টম কল বিবেচনা করুন.
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);
সিপিইউ এবং জিপিইউ উভয়েই, একটি টিপলকে পয়েন্টারগুলির একটি অ্যারে হিসাবে মেমরিতে উপস্থাপন করা হয়। 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 ইনপুট একটি সুবিধা, কিন্তু তারা কঠোরভাবে প্রয়োজনীয় নয়। যদি আমরা কাস্টম কলে টিপল ইনপুট সমর্থন না করি, তাহলে আপনি কাস্টম কলে পাঠানোর আগে গেট-টুপল-এলিমেন্ট ব্যবহার করে টিপলগুলিকে সর্বদা আনপ্যাক করতে পারেন।
অন্যদিকে, টিপল আউটপুট আপনাকে এমন কিছু করতে দেয় যা আপনি অন্যথায় করতে পারেন না।
টিপল আউটপুট থাকার সুস্পষ্ট কারণ হল টিপল আউটপুট হল কিভাবে একটি কাস্টম কল (বা অন্য কোন XLA অপ) একাধিক স্বাধীন অ্যারে ফেরত দেয়।
তবে কম স্পষ্টতই, একটি টিপল আউটপুট আপনার কাস্টম কল টেম্প মেমরি দেওয়ার একটি উপায়। হ্যাঁ, একটি আউটপুট একটি টেম্প বাফার প্রতিনিধিত্ব করতে পারে। বিবেচনা করুন, একটি আউটপুট বাফারের এমন বৈশিষ্ট্য রয়েছে যা op এটিতে লিখতে পারে এবং এটি লেখার পরে এটি থেকে পড়তে পারে। যে ঠিক কি আপনি একটি টেম্প বাফার থেকে চান.
উপরের উদাহরণে, ধরুন আমরা একটি টেম্প বাফার হিসাবে F32[1024] ব্যবহার করতে চাই। তারপরে আমরা উপরের মতো HLO লিখব এবং আমরা কাস্টম কলের আউটপুটের টিপল সূচক 1 কখনই পড়ব না।
CPU কাস্টম কলে Tuples
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