Trang này được dịch bởi Cloud Translation API.
Switch to English

Cuộc gọi tùy chỉnh XLA

Tài liệu này mô tả cách viết và sử dụng "lệnh gọi tùy chỉnh" XLA. Lệnh gọi tùy chỉnh cho phép bạn gọi mã được viết bằng ngôn ngữ lập trình như C ++ hoặc CUDA từ chương trình XLA.

Cuộc gọi tùy chỉnh trên CPU

Bạn có thể tạo một hướng dẫn HLO đại diện cho một cuộc gọi tùy chỉnh thông qua API ứng dụng khách của XLA. Điều này không được tiết lộ qua TensorFlow kể từ khi viết.

Ví dụ: đoạn mã sau sử dụng lệnh gọi tùy chỉnh để tính A[i] = B[i % 128] + C[i] trên CPU. (Tất nhiên bạn có thể - và nên làm! - làm điều này với HLO thông thường.)

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

Lưu ý rằng hàm do_custom_call cần biết kích thước của bộ đệm mà nó hoạt động. Trong ví dụ này, chúng tôi mã hóa các kích thước 128 và 2048. Nếu bạn không muốn làm điều này, bạn có thể chuyển các kích thước dưới dạng tham số cho lệnh gọi.

Cuộc gọi tùy chỉnh trên GPU

Khung cuộc gọi tùy chỉnh GPU hơi khác so với trên CPU. Đây là một ví dụ CUDA thực hiện tính toán A[i] = B[i % 128] + C[i] như mã CPU ở trên.

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

Trước tiên, hãy lưu ý rằng chức năng gọi tùy chỉnh GPU vẫn là một chức năng được thực thi trên CPU . Chức năng CPU do_custom_call của chúng tôi chịu trách nhiệm xếp hàng đợi công việc trên GPU. Ở đây, nó khởi chạy một nhân CUDA, nhưng nó cũng có thể làm một việc khác, chẳng hạn như gọi cublas.

buffers là một mảng các con trỏ sống trên máy chủ và mỗi phần tử mà nó chứa các điểm đến bộ nhớ thiết bị (tức là GPU). Các tham số đến trước, sau đó là giá trị đầu ra. Điều này đặc biệt khác với quy ước gọi CPU, có hai tham số, insout . Lý do chính mà chúng tôi phân kỳ là để có thể xử lý đầu vào / đầu ra hình tuple một cách hiệu quả; xem phần bên dưới.

Như trong ví dụ về CPU, chúng tôi đã mã hóa các kích thước bộ đệm đầu vào và đầu ra vào lệnh gọi tùy chỉnh của chúng tôi. Tuy nhiên, không giống như trong trường hợp CPU, việc chuyển các kích thước bộ đệm dưới dạng toán hạng cho lệnh gọi tùy chỉnh sẽ không hoạt động tốt. Thông thường, chúng ta cần các kích thước bộ đệm có sẵn cho chúng ta trên CPU; Ví dụ: khi khởi chạy kernel, chúng ta cần biết kích thước khối / lưới để sử dụng. Nhưng nếu chúng ta chuyển các kích thước bộ đệm dưới dạng toán hạng cho lệnh gọi tùy chỉnh của mình, các giá trị của chúng sẽ nằm trong bộ nhớ GPU. Sau đó, chúng tôi sẽ phải thực hiện một bản ghi nhớ đồng bộ thiết bị đến máy chủ đắt tiền khi bắt đầu hoạt động chỉ để đọc kích thước.

Để giúp bạn giải quyết vấn đề này, chúng tôi cung cấp thông số opaque . Bạn có thể đặt điều này thành một chuỗi byte tùy ý khi bạn tạo lệnh gọi tùy chỉnh:

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

xla::Shape có biểu diễn bộ đệm giao thức, bạn có thể lưu trữ proto tuần tự này bên trong opaque và giải mã hóa nó trong lệnh gọi tùy chỉnh GPU của bạn. Tuy nhiên, lưu ý rằng mặc dù xla::ShapeProto không thay đổi thường xuyên nhưng nó vẫn thay đổi. Kiểm tra nhật ký git để xem nó đã thay đổi như thế nào trong quá khứ.

Chuyển các bộ giá trị cho các cuộc gọi tùy chỉnh

Hãy xem xét cuộc gọi tùy chỉnh sau đây.

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

Trên cả CPU và GPU, một tuple được biểu diễn trong bộ nhớ dưới dạng một mảng con trỏ. Trong C ++ - mã giả, tham số 0 ở trên được trình bày như sau.

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

Mặc dù biểu diễn trong bộ nhớ của các bộ giá trị giống nhau trong CPU và GPU, nhưng chúng được xử lý khác nhau trong các quy ước gọi cuộc gọi tùy chỉnh của CPU và GPU.

Đầu ra Tuple dưới dạng bộ đệm tạm thời

Đầu vào Tuple cho các cuộc gọi tùy chỉnh là một sự tiện lợi, nhưng chúng không hoàn toàn cần thiết. Nếu chúng tôi không hỗ trợ đầu vào tuple cho các cuộc gọi tùy chỉnh, bạn luôn có thể giải nén các bộ giá trị bằng cách sử dụng get-tuple-element trước khi chuyển chúng đến cuộc gọi tùy chỉnh.

Mặt khác, đầu ra tuple cho phép bạn làm những điều bạn không thể.

Lý do rõ ràng để có đầu ra tuple là, đó là cách một lệnh gọi tùy chỉnh (hoặc bất kỳ op XLA nào khác) trả về nhiều mảng độc lập.

Nhưng ít rõ ràng hơn, đầu ra tuple cũng là một cách để cung cấp bộ nhớ tạm thời cuộc gọi tùy chỉnh của bạn. Có, một đầu ra có thể đại diện cho một bộ đệm tạm thời. Hãy xem xét, một bộ đệm đầu ra có thuộc tính mà op có thể ghi vào nó và nó có thể đọc từ nó sau khi được ghi vào. Đó chính xác là những gì bạn muốn từ một bộ đệm tạm thời.

Trong ví dụ trên, giả sử chúng ta muốn sử dụng F32[1024] làm bộ đệm tạm thời. Sau đó, chúng tôi sẽ viết HLO giống như ở trên và chúng tôi chỉ đơn giản là sẽ không bao giờ đọc chỉ mục tuple 1 của đầu ra của lệnh gọi tùy chỉnh.

Tuples trong cuộc gọi tùy chỉnh CPU

Trong mã CPU, chúng ta có một hàm do_custom_call(const void** ins, void* out) . ins là một mảng chỉ có một phần tử, trỏ đến param0 . Các bộ đệm con của param0 có thể truy cập được bằng cách bỏ tham chiếu con trỏ đó và các bộ đệm con của output_tuple có thể truy cập được bằng cách out tham chiếu.

Tuples trong cuộc gọi tùy chỉnh GPU

Trong mã GPU, chúng ta có một hàm do_custom_call(..., void** buffers, ...) . Trong trường hợp này, buffers là một mảng chủ gồm sáu con trỏ thiết bị, một con cho mỗi bộ đệm lá trong đầu vào / đầu ra. Để tạo danh sách phẳng, chúng tôi lặp lại các tham số và đầu ra, và đối với mỗi tham số, chúng tôi thực hiện việc duyệt qua hình dạng của nó. Cụ thể:

// 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