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

Tài liệu này mô tả cách viết và sử dụng các lệnh gọi tuỳ chỉnh XLA bằng thư viện XLA FFI. Lệnh gọi tuỳ chỉnh là một cơ chế để mô tả "thao tác" bên ngoài trong mô-đun HLO cho trình biên dịch XLA (tại thời gian biên dịch) và XLA FFI là một cơ chế để đăng ký việc triển khai các thao tác đó với XLA (tại thời điểm chạy). FFI là viết tắt của "giao diện hàm bên ngoài" và là một tập hợp các API C xác định giao diện nhị phân (ABI) cho XLA để gọi vào mã bên ngoài được viết bằng các ngôn ngữ lập trình khác. XLA cung cấp các liên kết chỉ có tiêu đề cho XLA FFI được viết bằng C++, giúp ẩn tất cả thông tin chi tiết cấp thấp của các API C cơ bản khỏi người dùng cuối.

JAX + XLA Custom Calls

Xem tài liệu JAX để biết các ví dụ từ đầu đến cuối về việc tích hợp các lệnh gọi tuỳ chỉnh và XLA FFI với JAX.

Liên kết FFI XLA

Liên kết XLA FFI là một quy cách thời gian biên dịch của chữ ký lệnh gọi tuỳ chỉnh: đối số lệnh gọi tuỳ chỉnh, thuộc tính và các loại của chúng, cũng như các tham số bổ sung được truyền qua ngữ cảnh thực thi (tức là luồng gpu cho phần phụ trợ GPU). Liên kết XLA FFI có thể được liên kết với mọi hàm có thể gọi C++ (con trỏ hàm, lambda, v.v.) có chữ ký operator() tương thích. Trình xử lý được tạo sẽ giải mã khung lệnh gọi XLA FFI (do API C ổn định xác định), kiểm tra kiểu của tất cả các tham số và chuyển tiếp kết quả đã giải mã đến lệnh gọi lại do người dùng xác định.

Liên kết FFI XLA phụ thuộc nhiều vào việc lập trình siêu mẫu để có thể biên dịch trình xử lý đã tạo thành mã máy hiệu quả nhất. Thời gian chạy có độ trễ theo thứ tự vài nano giây cho mỗi tham số lệnh gọi tuỳ chỉnh.

Các điểm tuỳ chỉnh XLA FFI được triển khai dưới dạng các chuyên môn hoá mẫu và người dùng có thể xác định cách giải mã các loại tuỳ chỉnh của họ, tức là có thể xác định quy trình giải mã tuỳ chỉnh cho các loại enum class do người dùng xác định.

Trả về lỗi từ các lệnh gọi tuỳ chỉnh

Các phương thức triển khai lệnh gọi tuỳ chỉnh phải trả về giá trị xla::ffi::Error để báo hiệu thành công hoặc lỗi cho thời gian chạy XLA. Nó tương tự như absl::Status và có cùng một tập hợp mã lỗi. Chúng tôi không sử dụng absl::Status vì nó không có ABI ổn định và sẽ không an toàn khi truyền nó giữa thư viện lệnh gọi tuỳ chỉnh được tải động và chính XLA.

// Handler that always returns an error.
auto always_error = Ffi::Bind().To(
    []() { return Error(ErrorCode::kInternal, "Oops!"); });

// Handler that always returns a success.
auto always_success = Ffi::Bind().To(
    []() { return Error::Success(); });

Đối số và kết quả của vùng đệm

XLA sử dụng kiểu truyền đích cho kết quả: các lệnh gọi tuỳ chỉnh (hoặc bất kỳ thao tác XLA nào khác) không phân bổ bộ nhớ cho kết quả mà thay vào đó, ghi vào các đích được truyền bởi thời gian chạy XLA. XLA sử dụng việc chỉ định vùng đệm tĩnh và phân bổ vùng đệm cho tất cả các giá trị dựa trên phạm vi hoạt động của chúng tại thời gian biên dịch.

Kết quả được truyền đến các trình xử lý FFI được gói trong một mẫu Result<T> có ngữ nghĩa giống như con trỏ: operator-> cho phép truy cập vào tham số cơ bản.

Đối số và kết quả AnyBuffer cho phép truy cập vào các tham số bộ đệm cuộc gọi tuỳ chỉnh của mọi loại dữ liệu. Điều này hữu ích khi lệnh gọi tuỳ chỉnh có một cách triển khai chung hoạt động cho nhiều loại dữ liệu và cách triển khai lệnh gọi tuỳ chỉnh thực hiện việc phân phối thời gian chạy dựa trên loại dữ liệu. AnyBuffer cho phép truy cập vào loại dữ liệu vùng đệm, các phương diện và một con trỏ trỏ đến chính vùng đệm.

%0 = "stablehlo.custom_call"(%arg0) {
  call_target_name = "foo",
  api_version = 4 : i32
} : (tensor<2x2xf32>) -> tensor<2x2xf32>

// Buffers of any number of dimensions and data type.
auto handler = Ffi::Bind().Arg<AnyBuffer>().Ret<AnyBuffer>().To(
    [](AnyBuffer arg, Result<AnyBuffer> res) -> Error {
      void* arg_data = arg.untyped_data();
      void* res_data = res->untyped_data();
      return Error::Success();
    });

Đối số và kết quả của vùng đệm bị hạn chế

Buffer cho phép thêm các ràng buộc vào kiểu dữ liệu vùng đệm và số lượng phương diện, đồng thời các ràng buộc này sẽ được trình xử lý tự động kiểm tra và trả về lỗi cho thời gian chạy XLA, nếu các đối số thời gian chạy không khớp với chữ ký trình xử lý FFI.

// Buffers of any number of dimensions and F32 data type.
auto handler = Ffi::Bind().Arg<Buffer<F32>>().Ret<Buffer<F32>>().To(
    [](Buffer<F32> arg, Result<Buffer<F32>> res) -> Error {
      float* arg_data = arg.typed_data();
      float* res_data = res->typed_data();
      return Error::Success();
    });

// Buffers of number of dimensions 2 and F32 data type.
auto handler = Ffi::Bind().Arg<BufferR2<F32>>().Ret<BufferR2<F32>>().To(
    [](BufferR2<F32> arg, Result<BufferR2<F32>> res) -> Error {
      float* arg_data = arg.typed_data();
      float* res_data = res->typed_data();
      return Error::Success();
    });

So khớp và xác minh vùng đệm

Các mẫu vùng đệm có thể thể hiện những quy tắc ràng buộc cụ thể hơn so với các loại đối số và kết quả trong một liên kết FFI. Chúng rất hữu ích để tinh chỉnh AnyBuffer thành một loại vùng đệm cụ thể, kiểm tra kích thước phương diện và kiểm tra mối quan hệ giữa nhiều hình dạng vùng đệm.

Các mẫu là bất biến và có thể chỉ định dtype và thứ hạng trực tiếp hoặc bằng các đối tượng sửa đổi tương ứng:

namespace m = ::xla::ffi::match;

m::Buffer<F32, 2>();
m::Buffer().WithDType<F32>().WithRank<2>();

Hàm Match kiểm tra một mẫu và trả về loại vùng đệm bị hạn chế tối đa của mẫu đó. Để tinh chỉnh một AnyBuffer, mẫu phải chỉ định chính xác một kiểu dữ liệu và một thứ hạng. Việc so khớp một vùng đệm đã nhập sẽ xác minh mọi ràng buộc bổ sung và giữ nguyên kiểu của vùng đệm đó.

namespace m = ::xla::ffi::match;

auto handler = Ffi::Bind().Arg<AnyBuffer>().Ret<AnyBuffer>().To(
    [](AnyBuffer input, Result<AnyBuffer> output) -> Error {
      int64_t rows;
      int64_t cols;

      ErrorOr<BufferR2<F32>> typed_input =
          Match("input", input,
                m::Buffer<F32>().WithDims(&rows, &cols));
      if (typed_input.has_error()) {
        return std::move(typed_input).error();
      }

      ErrorOr<Result<BufferR2<F32>>> typed_output =
          Match("output", output,
                m::Buffer<F32>().WithDims(rows, cols));
      if (typed_output.has_error()) {
        return std::move(typed_output).error();
      }

      float* input_data = typed_input->typed_data();
      float* output_data = (*typed_output)->typed_data();
      return Error::Success();
    });

Trong ví dụ này, input khớp sẽ ghi lại hai kích thước của phương diện và output khớp sẽ xác minh rằng phương diện đó có cùng hình dạng. Các giá trị được ghi lại chỉ được xác nhận khi mẫu hoàn chỉnh khớp thành công.

Mẫu kích thước có thể không bị ràng buộc, cố định hoặc được chụp. Các giá trị cố định và con trỏ chụp được chuyển đổi ngầm định thành các mẫu phương diện:

m::Dim()       // Any dimension size.
m::Dim(16)     // A dimension whose size is 16.
m::Dim(&rows)  // Any dimension size, captured in `rows` on success.

m::Buffer<F32>().WithDims(&rows, 16);

WithDims mô tả một hình dạng hoàn chỉnh và đặt thứ hạng của hình dạng đó dựa trên số lượng đối số. WithDim<I> hoặc WithDim(index, ...) hạn chế thứ nguyên theo vị trí mà không cố định thứ hạng hoàn chỉnh. Một ràng buộc về vị trí yêu cầu phải có phương diện được tham chiếu nhưng không giới hạn mọi phương diện khác:

ErrorOr<BufferR4<F32>> MatchActivations(AnyBuffer input) {
  return Match(
      "input", input,
      m::Buffer<F32, 4>().WithDim<0>(1).WithDim<3>(128));
}

Error VerifyDimension(AnyBuffer input, size_t index, int64_t size) {
  return Verify("input", input, m::Buffer().WithDim(index, size));
}

MatchActivations yêu cầu một vùng đệm F32 hạng 4 có kích thước 0 và 3 bằng 1 và 128, tương ứng; kích thước 1 và 2 không bị hạn chế. Nếu không có WithRank rõ ràng, WithDim(index, ...) sẽ chấp nhận mọi thứ hạng lớn hơn index.

Ảnh chụp kích thước cũng có thể thể hiện mối quan hệ giữa các vị trí:

auto matrix = m::Buffer<F32>().WithDims(&rows, 128);

int64_t n;
ErrorOr<BufferR2<F32>> square =
    Match("matrix", buffer,
          m::Buffer<F32>().WithDims(&n, &n));

Khi cùng một con trỏ chụp xuất hiện nhiều lần, tất cả các phương diện tương ứng phải có cùng kích thước. Do đó, lệnh gọi Match ở trên chỉ chấp nhận ma trận vuông. Hệ thống chỉ ghi lại thông tin về phương diện sau khi mẫu hoàn chỉnh thành công và không thay đổi khi thất bại.

Mối quan hệ toàn bộ vùng đệm không yêu cầu bạn phải nắm bắt mọi phương diện. WithShapeOf khớp với hình dạng thời gian chạy hoàn chỉnh của một vùng đệm khác trong khi cho phép một dtype khác. Like cũng yêu cầu cùng một dtype:

Error VerifySortBuffers(AnyBuffer keys, AnyBuffer values,
                        Result<AnyBuffer> keys_output) {
  Error error =
      Verify("values", values, m::Buffer().WithShapeOf(keys));
  if (error.failure()) {
    return error;
  }

  return Verify("keys output", *keys_output, m::Buffer().Like(keys));
}

Cả hai đối tượng sửa đổi đều sao chép siêu dữ liệu của vùng đệm tham chiếu khi mẫu được tạo; mẫu không giữ lại tham chiếu đến vùng đệm. Vì đây là các ràng buộc thời gian chạy, nên chúng không tinh chỉnh AnyBuffer thành một kiểu dữ liệu trả về cụ thể. Chúng chủ yếu hữu ích với Verify hoặc với Match khi vùng đệm đầu vào đã có một loại cụ thể.

Sử dụng Verify khi vùng đệm phải giữ nguyên loại hiện có hoặc khi một mẫu chấp nhận nhiều loại dữ liệu hoặc thứ hạng. Ví dụ: một vùng đệm chỉ mục có thể chấp nhận S32 hoặc S64 và gửi theo kiểu dữ liệu sau khi xác minh:

namespace m = ::xla::ffi::match;

Error VerifyIndices(AnyBuffer indices) {
  auto pattern = m::Buffer().WithDType<S32, S64>().WithRank<1, 2>();

  Error error = Verify("indices", indices, pattern);
  if (error.failure()) {
    return error;
  }

  // Dispatch an implementation based on `indices.element_type()`.
  return Error::Success();
}

Vì mẫu này cho phép nhiều loại vùng đệm cụ thể, nên mẫu này không thể tinh chỉnh AnyBuffer bằng Match: không có loại dữ liệu trả về duy nhất. Bạn vẫn có thể truyền đến Match bằng một vùng đệm đã nhập, kiểu dữ liệu trả về của vùng đệm này đã được xác định. Bạn cũng có thể xác minh một vùng đệm đã nhập; điều này thường hữu ích khi kiểm tra mối quan hệ hình dạng:

Error VerifySameShape(BufferR2<F32> lhs, BufferR2<F32> rhs) {
  int64_t rows;
  int64_t cols;

  Error error = Verify(
      "lhs", lhs,
      m::Buffer().WithDims(&rows, &cols));
  if (error.failure()) {
    return error;
  }

  return Verify("rhs", rhs,
                m::Buffer().WithDims(rows, cols));
}

Tên được truyền dưới dạng đối số đầu tiên đến Match và Verify là bắt buộc và được đưa vào lỗi, cùng với nội dung mô tả về ràng buộc không thành công.

API FFI bên ngoài trong xla/ffi/api/ffi.h trả về Error từ Verify và ErrorOr<T> từ Match, như trong các ví dụ trên. API FFI nội bộ trong xla/ffi/ffi.h cung cấp cùng một giao diện so khớp bằng cách sử dụng absl::Status và absl::StatusOr<T>:

ABSL_ASSIGN_OR_RETURN(BufferR2<F32> input,
                 Match("input", buffer, m::Buffer<F32, 2>()));
ABSL_RETURN_IF_ERROR(Verify(
    "input", input,
    m::Buffer().WithDims(expected_rows, m::Dim())));

Đối số và kết quả có thể thay đổi

Nếu số lượng đối số và kết quả có thể khác nhau trong các phiên bản khác nhau của một lệnh gọi tuỳ chỉnh, thì chúng có thể được giải mã tại thời gian chạy bằng cách sử dụng RemainingArgs và RemainingRets.

auto handler = Ffi::Bind().RemainingArgs().RemainingRets().To(
    [](RemainingArgs args, RemainingRets results) -> Error {
      ErrorOr<AnyBuffer> arg = args.get<AnyBuffer>(0);
      ErrorOr<Result<AnyBuffer>> res = results.get<AnyBuffer>(0);

      if (!arg.has_value()) {
        return Error(ErrorCode::kInternal, arg.error());
      }

      if (!res.has_value()) {
        return Error(ErrorCode::kInternal, res.error());
      }

      return Error::Success();
    });

Bạn có thể khai báo các đối số và kết quả có độ dài thay đổi sau các đối số và kết quả thông thường, tuy nhiên, việc liên kết các đối số và kết quả thông thường sau đối số có độ dài thay đổi là không hợp lệ.

auto handler =
    Ffi::Bind()
        .Arg<AnyBuffer>()
        .RemainingArgs()
        .Ret<AnyBuffer>()
        .RemainingRets()
        .To([](AnyBuffer arg, RemainingArgs args, AnyBuffer ret,
               RemainingRets results) -> Error { return Error::Success(); });

Thuộc tính

FFI XLA hỗ trợ việc tự động giải mã mlir::DictionaryAttr được truyền dưới dạng custom_call backend_config thành các đối số trình xử lý FFI.

%0 = "stablehlo.custom_call"(%arg0) {
  call_target_name = "foo",
  backend_config= {
    i32 = 42 : i32,
    str = "string"
  },
  api_version = 4 : i32
} : (tensor<f32>) -> tensor<f32>

Trong ví dụ này, lệnh gọi tuỳ chỉnh có một đối số vùng đệm và hai thuộc tính, đồng thời XLA FFI có thể tự động giải mã các đối số và thuộc tính này rồi chuyển đến hàm có thể gọi do người dùng xác định.

auto handler = Ffi::Bind()
  .Arg<BufferR0<F32>>()
  .Attr<int32_t>("i32")
  .Attr<std::string_view>("str")
  .To([](BufferR0<F32> buffer, int32_t i32, std::string_view str) {
    return Error::Success();
  });

Thuộc tính Enum do người dùng xác định

XLA FFI có thể tự động giải mã các thuộc tính MLIR tích phân thành các enum do người dùng xác định. Lớp Enum phải có cùng kiểu số nguyên cơ bản và quá trình giải mã phải được đăng ký rõ ràng với XLA FFI.

%0 = "stablehlo.custom_call"(%arg0) {
  call_target_name = "foo",
  backend_config= {
    command = 0 : i32
  },
  api_version = 4 : i32
} : (tensor<f32>) -> tensor<f32>

enum class Command : int32_t {
  kAdd = 0,
  kMul = 1,
};

XLA_FFI_REGISTER_ENUM_ATTR_DECODING(Command);

auto handler = Ffi::Bind().Attr<Command>("command").To(
    [](Command command) -> Error { return Error::Success(); });

Liên kết tất cả thuộc tính tuỳ chỉnh của cuộc gọi

Bạn có thể truy cập vào tất cả các thuộc tính cuộc gọi tuỳ chỉnh dưới dạng từ điển và chỉ giải mã các thuộc tính cần thiết tại thời gian chạy.

auto handler = Ffi::Bind().Attrs().To([](Dictionary attrs) -> Error {
  ErrorOr<int32_t> i32 = attrs.get<int32_t>("i32");
  return Error::Success();
});

Thuộc tính Struct do người dùng xác định

FFI XLA có thể giải mã các thuộc tính từ điển thành các cấu trúc do người dùng xác định.

%0 = "stablehlo.custom_call"(%arg0) {
  call_target_name = "foo",
  backend_config= {
    range = { lo = 0 : i64, hi = 42 : i64 }
  },
  api_version = 4 : i32
} : (tensor<f32>) -> tensor<f32>

Trong ví dụ trên, range là một thuộc tính mlir::DictionaryAttr và thay vì truy cập vào các trường từ điển theo tên, bạn có thể tự động giải mã thuộc tính này dưới dạng một cấu trúc C++. Bạn phải đăng ký rõ ràng quá trình giải mã bằng macro XLA_FFI_REGISTER_STRUCT_ATTR_DECODING (ở phía sau, macro này sẽ xác định một chuyên môn hoá mẫu trong không gian tên ::xla::ffi, do đó, bạn phải thêm macro vào không gian tên chung).

struct Range {
  int64_t lo;
  int64_t hi;
};

XLA_FFI_REGISTER_STRUCT_ATTR_DECODING(Range, StructMember<int64_t>("lo"),
                                             StructMember<int64_t>("hi"));

auto handler = Ffi::Bind().Attr<Range>("range").To([](Range range) -> Error{
  return Error::Success();
});

Bạn có thể tải các thuộc tính tuỳ chỉnh từ một từ điển, giống như mọi thuộc tính khác. Trong ví dụ bên dưới, tất cả các thuộc tính cuộc gọi tuỳ chỉnh được giải mã dưới dạng Dictionary và bạn có thể truy cập vào range theo tên.

auto handler = Ffi::Bind().Attrs().To([](Dictionary attrs) -> Error {
  ErrorOr<Range> range = attrs.get<Range>("range");
  return Error::Success();
});

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

Bạn có thể tạo một chỉ dẫn HLO đại diện cho một lệnh gọi tuỳ chỉnh thông qua API máy khách của XLA. Ví dụ: đoạn mã sau đây sử dụng một 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 với 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}),
        /*opaque=*/"", /*has_side_effect=*/false,
        /*output_operand_aliasing=*/{}, /*literal=*/nullptr,
        /*schedule=*/CustomCallSchedule::SCHEDULE_NONE,
        /*api_version=*/CustomCallApiVersion::API_VERSION_TYPED_FFI);
}

// Constrain custom call arguments to 1-dimensional buffers of F32 data type.
using BufferF32 = xla::ffi::BufferR1<xla::ffi::DataType::F32>;

// Implement a custom call as a C++ function. Note that we can use `Buffer` type
// defined by XLA FFI that gives us access to buffer data type and shape.
xla::ffi::Error do_custom_call(BufferF32 in0, BufferF32 in1,
                               xla::ffi::Result<BufferF32> out) {
  size_t d0 = in0.dimensions[0];
  size_t d1 = in1.dimensions[0];

  // Check that dimensions are compatible.
  assert(out->dimensions[0] == d1 && "unexpected dimensions");

  for (size_t i = 0; i < d1; ++i) {
    out->data[i] = in0.data[i % d0] + in1.data[i];
  }
}

// Explicitly define an XLA FFI handler signature and bind it to the
// `do_custom_call` implementation. XLA FFI handler can automatically infer
// type signature from the custom call function, but it relies on magical
// template metaprogramming an explicit binding provides and extra level of
// type checking and clearly states custom call author intentions.
XLA_FFI_DEFINE_HANDLER(handler, do_custom_call,
                       ffi::Ffi::Bind()
                           .Arg<Buffer>()
                           .Arg<Buffer>()
                           .Ret<Buffer>());

// Registers `handler` with and XLA FFI on a "Host" platform.
XLA_FFI_REGISTER_HANDLER(xla::ffi::GetXlaFfiApi(), "do_custom_call",
                         "Host", handler);

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

Việc đăng ký lệnh gọi tuỳ chỉnh GPU bằng XLA FFI gần như giống hệt nhau, điểm khác biệt duy nhất là đối với GPU, bạn cần yêu cầu một luồng nền tảng cơ bản (luồng CUDA hoặc ROCM) để có thể khởi chạy hạt nhân trên thiết bị. Sau đây là một ví dụ về CUDA thực hiện cùng một phép tính (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, BufferF32 in0, BufferF32 in1,
                    xla::ffi::Result<BufferF32> out) {
  size_t d0 = in0.dimensions[0];
  size_t d1 = in1.dimensions[0];
  size_t d2 = out->dimensions[0];

  assert(d0 == 128 && d1 == 2048 && d2 == 2048 && "unexpected dimensions");

  const int64_t block_dim = 64;
  const int64_t grid_dim = 2048 / block_dim;
  custom_call_kernel<<<grid_dim, block_dim, 0, stream>>>(
    in0.data, in1.data, out->data);
}

XLA_FFI_DEFINE_HANDLER(handler, do_custom_call,
                       ffi::Ffi::Bind()
                           .Ctx<xla::ffi::PlatformStream<CUstream>>()
                           .Arg<BufferF32>()
                           .Arg<BufferF32>()
                           .Ret<BufferF32>());

XLA_FFI_REGISTER_HANDLER(xla::ffi::GetXlaFfiApi(), "do_custom_call",
                         "CUDA", handler);

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

Các đối số và kết quả cũng nằm trên máy chủ lưu trữ, đồng thời thành phần dữ liệu chứa một con trỏ đến bộ nhớ thiết bị (tức là GPU). Các vùng đệm được truyền đến trình xử lý lệnh gọi tuỳ chỉnh có hình dạng của các vùng đệm thiết bị cơ bản, vì vậy, lệnh gọi tuỳ chỉnh có thể tính toán các tham số khởi chạy nhân từ các vùng đệm đó.

Truyền các bộ đến các lệnh gọi tuỳ chỉnh

Hãy xem xét lệnh gọi tuỳ chỉnh sau.

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ỏ. Khi XLA gọi các lệnh gọi tuỳ chỉnh bằng đối số hoặc kết quả của bộ giá trị, XLA sẽ làm phẳng các đối số hoặc kết quả đó và truyền dưới dạng đối số hoặc kết quả của vùng đệm thông thường.

Đầu ra của bộ dữ liệu dưới dạng vùng đệm tạm thời

Đầu vào của bộ giá trị cho các lệnh gọi tuỳ chỉnh là một tiện ích, nhưng không thực sự cần thiết. Nếu chúng tôi không hỗ trợ các đầu vào bộ giá trị cho các lệnh gọi tuỳ 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 truyền chúng đến lệnh gọi tuỳ chỉnh.

Mặt khác, đầu ra của bộ giá trị cho phép bạn làm những việc mà bạn không thể làm được nếu không có bộ giá trị.

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

Tuy nhiên, ít rõ ràng hơn, đầu ra của bộ giá trị 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ủa bạn. Có, đầu ra có thể biểu thị một vùng đệm tạm thời. Hãy xem xét, một vùng đệm đầu ra có thuộc tính mà thao tác có thể ghi vào đó và có thể đọc từ đó sau khi được ghi vào. Đó chính xác là những gì bạn muốn ở mộ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 giống như ở trên và chúng ta sẽ không bao giờ đọc chỉ mục bộ dữ liệu 1 của đầu ra lệnh gọi tuỳ chỉnh.