Lệnh gọi tuỳ chỉnh XLA

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

Tạo lệnh gọi tuỳ chỉnh trên CPU

Bạn có thể tạo hướng dẫn HLO biểu thị lệnh gọi tuỳ chỉnh thông qua API ứng dụng khách của XLA. Ví dụ: Mã sau đây sử dụng lệnh gọi tuỳ chỉnh để tính toán A[i] = B[i % 128]+ C[i] trên CPU. (Tất nhiên là bạn có thể – và nên làm! – thực hiện việc này bằng HLO thông thường.)

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

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

Tạo lệnh gọi tuỳ chỉnh trên GPU

Khung lệnh gọi tuỳ chỉnh trong GPU có chút khác biệt so với khung trên CPU. Dưới đây là ví dụ về CUDA thực hiện phép tính tương tự (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 = 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");

Trước tiên, hãy lưu ý rằng hàm gọi hàm tuỳ chỉnh GPU vẫn là một hàm được thực thi trên CPU. Hàm CPU do_custom_call chịu trách nhiệm đưa công việc vào hàng đợi trên GPU. Ở đây, mã này sẽ khởi chạy một hạt nhân CUDA, nhưng cũng có thể thực hiện một tác vụ khác, chẳng hạn như gọi cuBLAS.

buffers là một mảng con trỏ nằm trên máy chủ lưu trữ và mỗi phần tử trong đó đều chứa các điểm đến bộ nhớ thiết bị (cụ thể là GPU). Các tham số xuất hiện trước, theo sau là giá trị đầu ra. Điều này khác biệt đáng kể so với quy ước gọi CPU, có hai tham số là ins và out. Quy ước gọi GPU giúp bạn có thể xử lý các đầu vào/đầu ra có dạng bộ dữ liệu một cách hiệu quả.

Như trong ví dụ về CPU, chúng tôi đã mã hoá cứng kích thước bộ nhớ đệm đầu vào và đầu ra vào lệnh gọi tuỳ chỉnh. Tuy nhiên, không giống như trong trường hợp CPU, việc chuyển dung lượng bộ nhớ đệm dưới dạng hoạt động đến lệnh gọi tuỳ chỉnh sẽ không hoạt động tốt. Thông thường, chúng ta cần có dung lượng bộ nhớ đệm trên CPU (ví dụ: khi khởi chạy hạt nhân, 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 kích thước bộ nhớ đệm dưới dạng toán tử vào lệnh gọi tuỳ chỉnh, thì giá trị của chúng sẽ nằm trong bộ nhớ GPU. Sau đó, chúng ta sẽ phải thực hiện một memcpy đồng bộ – với máy chủ lưu trữ tốn kém khi bắt đầu thao tác 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 tham số opaque. Bạn có thể đặt giá trị này thành một chuỗi byte tuỳ ý khi tạo lệnh gọi tuỳ chỉnh:

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

Vì xla::Shape có bản trình bày vùng đệm giao thức, nên bạn có thể lưu trữ proto chuyển đổi tuần tự này bên trong opaque và giải tuần tự nó trong lệnh gọi tuỳ chỉnh GPU. Tuy nhiên, lưu ý rằng mặc dù xla::ShapeProto không thay đổi thường xuyên, nhưng mã này sẽ thay đổi. Kiểm tra nhật ký Git để xem nhật ký này đã thay đổi như thế nào trước đây.

Báo hiệu có lỗi

Nếu lệnh gọi tuỳ chỉnh của bạn gặp lỗi, bạn có thể báo hiệu lỗi đó đến môi trường thời gian chạy XLA (thay vì sự cố hoặc trả về dữ liệu vô nghĩa trong vùng đệm đầu ra) bằng cách sử dụng chữ ký sau cho hàm của bạn:

Trên CPU:

#include "xla/service/custom_call_status.h"

void do_custom_call(void* out, const void** in, XlaCustomCallStatus* status);

trên 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);

Bạn có thể báo hiệu lỗi bằng cách sử dụng XlaCustomCallStatusSetFailure, ví dụ:

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.
}

Bạn cũng có thể sử dụng XlaCustomCallStatusSetSuccess để chỉ báo trạng thái thành công, nhưng XlaCustomCallStatus đang ở trạng thái thành công theo mặc định, vì vậy, việc bỏ qua hoàn toàn thành công cũng sẽ cho biết trạng thái thành công.

Khi sử dụng hàm gọi tuỳ chỉnh với chữ ký này, bạn phải tạo hoạt động custom-call tương ứng với nhóm phiên bản API thích hợp, ví dụ:

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

Khi không thành công, sẽ không có dữ liệu đầu ra của lệnh gọi tuỳ chỉnh nào được sử dụng; thời gian chạy XLA sẽ kết thúc quá trình tính toán. Tính toán HLO không thể khôi phục lỗi (ví dụ: bằng cách phát hiện và xử lý lỗi).

Truyền bộ dữ liệu (tuples) đến các lệnh gọi tuỳ chỉnh

Hãy cân nhắc lệnh gọi tuỳ chỉnh sau đây.

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

Trên cả CPU và GPU, một bộ dữ liệu được biểu thị trong bộ nhớ dưới dạng một mảng con trỏ. Trong mã giả C++, 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ù cách trình bày bộ dữ liệu trong CPU và GPU là giống nhau trong bộ nhớ, nhưng chúng được xử lý khác nhau trong quy ước gọi lệnh gọi tuỳ chỉnh của CPU và GPU.

Chuyển đầu ra dưới dạng vùng đệm tạm thời

Việc thu thập dữ liệu đầu vào cho lệnh gọi tuỳ chỉnh rất tiện lợi, nhưng việc này là không cần thiết. Nếu chúng tôi không hỗ trợ bộ dữ liệu đầu vào cho lệnh gọi tuỳ chỉnh, thì bạn luôn có thể giải nén các bộ dữ liệu bằng cách sử dụng phần tử get-tuple-tuple trước khi chuyển các bộ dữ liệu đó đến lệnh gọi tuỳ chỉnh.

Mặt khác, đầu ra của bộ công cụ tuple cho phép bạn thực hiện những việc mà bạn không thể làm.

Lý do rõ ràng để có kết quả của bộ dữ liệu đầu ra là cách một lệnh gọi tuỳ chỉnh (hoặc bất kỳ hoạt động XLA nào khác) trả về nhiều mảng độc lập.

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

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

Lưu lượng truy cập trong lệnh gọi tuỳ chỉnh bằng 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. Bạn có thể truy cập vào các vùng đệm phụ của param0 bằng cách loại bỏ ưu tiên con trỏ đó, còn các vùng đệm phụ của output_tuple có thể truy cập được bằng cách tham chiếu out.

Bộ dữ liệu trong lệnh gọi tuỳ 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 máy chủ gồm 6 con trỏ thiết bị, một mảng cho mỗi vùng đệm lá trong đầu vào/đầu ra. Để tạo danh sách phẳng, chúng ta lặp lại các tham số và đầu ra. Đối với mỗi tham số, chúng ta sẽ truyền tải trước theo hình dạng của danh sách. 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