इस पेज का अनुवाद Cloud Translation API से किया गया है.
Switch to English

XLA कस्टम कॉल

यह दस्तावेज़ बताता है कि एक्सएलए "कस्टम कॉल" कैसे लिखना और उपयोग करना है। कस्टम कॉल आपको एक प्रोग्रामिंग भाषा में लिखे कोड को C ++ या CUDA जैसे XLA प्रोग्राम से लिखने की अनुमति देता है।

सीपीयू पर कस्टम-कॉल

आप एक HLO निर्देश बना सकते हैं जो XLA के क्लाइंट API के माध्यम से कस्टम-कॉल का प्रतिनिधित्व करता है। यह TensorFlow लेखन के रूप में उजागर नहीं है।

उदाहरण के लिए, निम्न कोड CPU पर 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 सीपीयू फ़ंक्शन GPU पर काम को रोकने के लिए जिम्मेदार है। यहाँ यह CUDA कर्नेल लॉन्च करता है, लेकिन यह कॉल क्यूबल्स की तरह कुछ और भी कर सकता है।

buffers पॉइंटर्स की एक सरणी है जो होस्ट पर रहता है, और प्रत्येक तत्व में डिवाइस (यानी GPU) मेमोरी के अंक शामिल हैं। पैरामीटर पहले आते हैं, उसके बाद आउटपुट वैल्यू आती है। यह सीपीयू कॉलिंग कन्वेंशन से विशेष रूप से अलग है, जिसमें दो परम, ins और out । मुख्य कारण जो हम कहते हैं कि यह टपल-आकार के इनपुट / आउटपुट को कुशलतापूर्वक संभालना संभव बनाता है; नीचे अनुभाग देखें।

सीपीयू उदाहरण के रूप में, हमने अपने कस्टम कॉल में इनपुट और आउटपुट बफर साइज को हार्डकोड किया है। हालांकि सीपीयू मामले के विपरीत, बफर कॉल को कस्टम कॉल के लिए ऑपरेंड के रूप में पास करना अच्छी तरह से काम नहीं करेगा। आमतौर पर हमें सीपीयू पर हमारे लिए उपलब्ध बफर आकारों की आवश्यकता होती है; उदाहरण के लिए, कर्नेल लॉन्च करते समय, हमें उपयोग करने के लिए ब्लॉक / ग्रिड आयामों को जानना होगा। लेकिन अगर हम बफर साइज को अपने कस्टम कॉल के ऑपरेंड के रूप में पास करते हैं, तो उनके मूल्य जीपीयू मेमोरी में रहते हैं। फिर हमें आकार को पढ़ने के लिए हमारे ऑपरेशन की शुरुआत में एक महंगा सिंक्रोनस डिवाइस-टू-होस्ट मेमसीपी करना होगा।

आपको इसके आसपास काम करने देने के लिए, हम opaque पैरामीटर प्रदान करते हैं। जब आप कस्टम कॉल बनाते हैं, तो आप इसे बाइट्स के मनमाने तरीके से सेट कर सकते हैं:

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

चूँकि xla::Shape में एक प्रोटोकॉल बफर प्रतिनिधित्व है, आप इस धारावाहिक xla::Shape को opaque अंदर संग्रहीत कर सकते हैं और इसे अपने GPU कस्टम-कॉल के भीतर दे सकते हैं। नोट तथापि कि हालांकि xla::ShapeProto अक्सर परिवर्तन नहीं करता है, यह बदल जाती है। यह कैसे अतीत में बदल गया है यह देखने के लिए गिट लॉग की जाँच करें।

कस्टम-कॉल को ट्यूपल्स पास करना

निम्नलिखित कस्टम-कॉल पर विचार करें।

 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;
 

यद्यपि सीपीयू और जीपीयू में टुपल्स का इन-मेमोरी प्रतिनिधित्व समान है, फिर भी उन्हें सीपीयू और जीपीयू कस्टम-कॉल कॉलिंग सम्मेलनों में अलग-अलग तरीके से नियंत्रित किया जाता है।

टपल आउटपुट टेम्पोरर बफ़र्स के रूप में

कस्टम-कॉल के लिए इनपुट इनपुट एक सुविधा है, लेकिन वे कड़ाई से आवश्यक नहीं हैं। अगर हमने कस्टम कॉल में टुपल इनपुट का समर्थन नहीं किया है, तो आप कस्टम कॉल पर जाने से पहले हमेशा ट्यूपल-गेट-टपल-एलिमेंट का उपयोग करके अनपैक कर सकते हैं।

दूसरी ओर, टपल आउटपुट आपको उन चीजों को करने देता है जो आप अन्यथा नहीं कर सकते थे।

स्पष्ट आउटपुट होने का स्पष्ट कारण यह है कि एक कस्टम कॉल (या कोई अन्य XLA op) कैसे कई स्वतंत्र सरणियों को लौटाता है।

लेकिन कम स्पष्ट रूप से, एक टपल आउटपुट भी आपके कस्टम कॉल टेम्प मेमोरी को देने का एक तरीका है। हां, एक आउटपुट एक अस्थायी बफर का प्रतिनिधित्व कर सकता है। गौर कीजिए, एक आउटपुट बफर में वह गुण होता है जो ऑप उस पर लिख सकता है, और इसे लिखे जाने के बाद इसे पढ़ सकता है। ठीक यही आप एक अस्थायी बफर से चाहते हैं।

ऊपर दिए गए उदाहरण में, मान लीजिए कि हम F32[1024] का उपयोग एक अस्थायी बफर के रूप में करना चाहते थे। तब हम HLO को ऊपर की तरह ही लिखेंगे, और हम कस्टम कॉल के आउटपुट के टपल इंडेक्स 1 को कभी नहीं पढ़ेंगे।

सीपीयू कस्टम-कॉल में ट्यूपल

CPU कोड में, हमारे पास एक फ़ंक्शन do_custom_call(const void** ins, void* out)ins सिर्फ एक तत्व के साथ एक सरणी है, जो param0 ओर param0 । की subbuffers param0 कि सूचक dereferencing से सुलभ हैं, और की subbuffers 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