本頁面由 Cloud Translation API 翻譯而成。
Switch to English

XLA自定義呼叫

本文檔介紹瞭如何編寫和使用XLA“自定義調用”。使用自定義調用,您可以從XLA程序調用以C ++或CUDA之類的編程語言編寫的代碼。

CPU上的自定義調用

您可以通過XLA的客戶端API創建代表自定義調用的HLO指令。撰寫本文時,這尚未通過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自定義調用框架與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 = 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自定義調用函數仍然是在CPU上執行的函數 。我們的do_custom_call CPU函數負責使GPU上的工作入隊。它在這裡啟動CUDA內核,但它也可以做其他事情,例如調用cublas。

buffers是駐留在主機上的指針數組,指針的每個元素都包含指向設備(即GPU)內存的指針。首先是參數,然後是輸出值。這與CPU調用約定有顯著不同,後者有兩個參數insout 。我們分歧的主要原因是可以有效地處理元組形的輸入/輸出。請參閱以下部分。

就像在CPU示例中一樣,我們已經將輸入和輸出緩衝區的大小硬編碼到我們的自定義調用中。但是,與CPU情況不同的是,將緩衝區大小作為操作數傳遞給自定義調用將無法正常工作。通常,我們需要CPU上可用的緩衝區大小。例如,在啟動內核時,我們需要知道要使用的塊/網格尺寸。但是,如果我們將緩衝區大小作為操作數傳遞給我們的自定義調用,則它們的值將存在於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內部,並在GPU定制調用中反序列化它。但是請注意,儘管xla::ShapeProto不會經常更改,但確實會更改。檢查git日誌以查看過去的變化。

將元組傳遞給自定義調用

考慮以下自定義調用。

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

在CPU和GPU上,元組在內存中均表示為指針數組。在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和GPU中相同,但是在CPU和GPU自定義調用調用約定中對元組的處理方式不同。

元組輸出作為臨時緩衝區

自定義調用的元組輸入很方便,但並非絕對必要。如果我們不支持將元組輸入到自定義調用中,則可以始終在將它們傳遞到自定義調用之前使用get-tuple-element打開元組的包裝。

另一方面,元組輸出可以讓您做其他事情無法做的事情。

具有元組輸出的明顯原因是,這就是自定義調用(或任何其他XLA op)返回多個獨立數組的方式。

但不太明顯的是,元組輸出也是提供自定義調用臨時內存的一種方法。是的, 輸出可以代表一個臨時緩衝區。考慮一下,輸出緩衝區具有op可以向其寫入的屬性,並且可以在寫入後對其進行讀取。這正是您想要的臨時緩衝區。

在上面的示例中,假設我們要使用F32[1024]作為臨時緩衝區。然後,我們將像上面那樣編寫HLO,並且根本就不會讀取自定義調用輸出的元組索引1。

CPU自定義調用中的元組

在CPU代碼中,我們有一個函數do_custom_call(const void** ins, void* out)ins是一個只有一個元素的數組,它指向param0 。通過解引用該指針可以訪問param0的子緩衝區,通過解引用out可以訪問output_tuple的子緩衝區。

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