تماس های سفارشی XLA

این سند نحوه نوشتن و استفاده از فراخوانی‌های سفارشی XLA را با استفاده از کتابخانه XLA FFI شرح می‌دهد. فراخوانی سفارشی مکانیزمی برای توصیف یک "عملیات" خارجی در ماژول HLO به کامپایلر XLA (در زمان کامپایل) است و XLA FFI مکانیزمی برای ثبت پیاده‌سازی چنین عملیاتی با XLA (در زمان اجرا) است. FFI مخفف "رابط تابع خارجی" است و مجموعه‌ای از APIهای C است که یک رابط دودویی (ABI) را برای XLA تعریف می‌کند تا به کد خارجی نوشته شده به زبان‌های برنامه‌نویسی دیگر فراخوانی شود. XLA اتصالات فقط سرصفحه‌ای را برای XLA FFI نوشته شده به زبان C++ فراهم می‌کند که تمام جزئیات سطح پایین APIهای C زیرین را از کاربر نهایی پنهان می‌کند.

تماس‌های سفارشی JAX + XLA

برای مثال‌های سرتاسری از ادغام فراخوانی‌های سفارشی و XLA FFI با JAX، به مستندات JAX مراجعه کنید.

اتصال XLA FFI

اتصال XLA FFI یک مشخصه زمان کامپایل از امضای فراخوانی سفارشی است: آرگومان‌های فراخوانی سفارشی، ویژگی‌ها و انواع آنها، و پارامترهای اضافی که از طریق زمینه اجرا (یعنی جریان gpu برای GPU backend) منتقل می‌شوند. اتصال XLA FFI می‌تواند به هر فراخوانی‌پذیر C++ (اشاره‌گر تابع، لامبدا و غیره) با امضای operator() سازگار متصل شود. کنترل‌کننده ساخته‌شده، قاب فراخوانی XLA FFI (تعریف‌شده توسط API پایدار C) را رمزگشایی می‌کند، تمام پارامترها را از نظر نوع بررسی می‌کند و نتایج رمزگشایی‌شده را به فراخوانی برگشتی تعریف‌شده توسط کاربر ارسال می‌کند.

اتصال XLA FFI به شدت به برنامه‌نویسی متای الگو متکی است تا بتواند هندلر ساخته شده را به کارآمدترین کد ماشین کامپایل کند. سربارهای زمان اجرا برای هر پارامتر فراخوانی سفارشی در حدود چند نانوثانیه است.

نقاط سفارشی‌سازی XLA FFI به عنوان تخصص‌های قالب پیاده‌سازی شده‌اند و کاربران می‌توانند نحوه رمزگشایی انواع سفارشی خود را تعریف کنند، یعنی می‌توان رمزگشایی سفارشی را برای انواع enum class تعریف شده توسط کاربر تعریف کرد.

برگرداندن خطاها از فراخوانی‌های سفارشی

پیاده‌سازی‌های فراخوانی سفارشی باید مقدار xla::ffi::Error را برای اعلام موفقیت یا خطا به زمان اجرای XLA برگردانند. این شبیه به absl::Status است و همان مجموعه کدهای خطا را دارد. ما absl::Status استفاده نمی‌کنیم زیرا ABI پایدار ندارد و انتقال آن بین کتابخانه فراخوانی سفارشی که به صورت پویا بارگذاری می‌شود و خود 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(); });

آرگومان‌ها و نتایج بافر

XLA از سبک ارسال به مقصد برای نتایج استفاده می‌کند: فراخوانی‌های سفارشی (یا هر عملیات XLA دیگری) حافظه‌ای برای نتایج اختصاص نمی‌دهند و در عوض در مقصدهای ارسالی توسط زمان اجرای XLA می‌نویسند. XLA از انتساب بافر استاتیک استفاده می‌کند و بافرها را برای همه مقادیر بر اساس محدوده‌های زنده آنها در زمان کامپایل اختصاص می‌دهد.

نتایج ارسالی به FFI handlerها در قالب Result<T> قرار می‌گیرند که معنایی شبیه به اشاره‌گر دارد: operator-> به پارامتر اصلی دسترسی می‌دهد.

آرگومان‌ها و نتایج AnyBuffer امکان دسترسی به پارامترهای بافر فراخوانی سفارشی از هر نوع داده‌ای را فراهم می‌کند. این زمانی مفید است که فراخوانی سفارشی دارای یک پیاده‌سازی عمومی باشد که برای چندین نوع داده کار می‌کند و پیاده‌سازی فراخوانی سفارشی، توزیع زمان اجرا را بر اساس نوع داده انجام می‌دهد. AnyBuffer امکان دسترسی به نوع داده بافر، ابعاد و یک اشاره‌گر به خود بافر را فراهم می‌کند.

%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();
    });

آرگومان‌ها و نتایج بافر محدود

Buffer امکان اضافه کردن محدودیت روی نوع داده بافر و تعداد ابعاد را فراهم می‌کند و اگر آرگومان‌های زمان اجرا با امضای FFI handler مطابقت نداشته باشند، به طور خودکار توسط handler بررسی می‌شوند و خطایی به XLA runtime برمی‌گردانند.

// 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();
    });

تطبیق و تأیید بافرها

الگوهای بافر می‌توانند محدودیت‌هایی را بیان کنند که خاص‌تر از انواع آرگومان و نتیجه در یک اتصال FFI هستند. آن‌ها برای اصلاح AnyBuffer به یک نوع بافر مشخص، بررسی اندازه ابعاد و بررسی روابط بین اشکال مختلف بافر مفید هستند.

الگوها تغییرناپذیر هستند و می‌توانند dtype و rank را مستقیماً یا با استفاده از اصلاح‌کننده‌های مربوطه مشخص کنند:

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

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

تابع Match یک الگو را بررسی می‌کند و نوع بافر با حداکثر محدودیت آن را برمی‌گرداند. برای اصلاح AnyBuffer ، الگو باید دقیقاً یک نوع داده و یک رتبه را مشخص کند. تطبیق بافر از قبل تایپ شده، هرگونه محدودیت اضافی را تأیید می‌کند و نوع آن را حفظ می‌کند.

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

در این مثال input تطبیق، اندازه‌های دو بُعدی آن را ثبت می‌کند و output تطبیق، یکسان بودن شکل آن را تأیید می‌کند. ثبت‌ها تنها زمانی انجام می‌شوند که الگوی کامل با موفقیت مطابقت داشته باشد.

یک الگوی بُعد می‌تواند نامحدود، ثابت یا ثبت‌شده باشد. مقادیر ثابت و اشاره‌گرهای ثبت‌شده به‌طور ضمنی به الگوهای بُعد تبدیل می‌شوند:

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 یک شکل کامل را توصیف می‌کند و رتبه آن را از تعداد آرگومان‌ها تعیین می‌کند. WithDim<I> یا WithDim(index, ...) ابعاد را از نظر موقعیتی محدود می‌کند بدون اینکه رتبه کامل را ثابت کند. یک محدودیت موقعیتی مستلزم وجود بُعد ارجاع‌شده است اما سایر ابعاد را بدون محدودیت باقی می‌گذارد:

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 به یک بافر F32 با درجه rank-4 با ابعاد 0 و 3 برابر با 1 و 128 نیاز دارد؛ ابعاد 1 و 2 بدون محدودیت هستند. بدون WithRank صریح، WithDim(index, ...) هر رتبه‌ای بزرگتر از index را می‌پذیرد.

ثبت ابعاد همچنین می‌تواند روابط بین موقعیت‌ها را بیان کند:

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

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

وقتی یک اشاره‌گر ثبت چندین بار ظاهر می‌شود، تمام ابعاد مربوطه باید اندازه یکسانی داشته باشند. بنابراین، فراخوانی Match در بالا فقط ماتریس‌های مربعی را می‌پذیرد. ثبت ابعاد فقط پس از موفقیت الگوی کامل نوشته می‌شوند و در صورت شکست بدون تغییر باقی می‌مانند.

روابط کل بافر نیازی به ثبت هر بُعد ندارند. WithShapeOf شکل کامل زمان اجرای بافر دیگری را تطبیق می‌دهد در حالی که اجازه می‌دهد نوع داده متفاوتی داشته باشد. Like نیز به همان نوع داده نیاز دارد:

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

هر دو اصلاح‌کننده، فراداده‌های بافر مرجع را هنگام ساخت الگو کپی می‌کنند؛ الگو ارجاعی به بافر را حفظ نمی‌کند. از آنجا که این‌ها محدودیت‌های زمان اجرا هستند، AnyBuffer را به یک نوع بازگشتی مشخص اصلاح نمی‌کنند. آن‌ها در درجه اول با Verify یا با Match زمانی که بافر ورودی از قبل یک نوع مشخص دارد، مفید هستند.

Verify زمانی استفاده کنید که بافر باید نوع موجود خود را حفظ کند، یا زمانی که یک الگو بیش از یک نوع داده یا رتبه را می‌پذیرد. برای مثال، یک بافر شاخص می‌تواند S32 یا S64 را بپذیرد و پس از تأیید، نوع داده خود را ارسال کند:

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

از آنجا که این الگو چندین نوع بافر مشخص را مجاز می‌داند، نمی‌تواند AnyBuffer با Match اصلاح کند: هیچ نوع بازگشتی منحصر به فردی وجود ندارد. همچنان می‌توان آن را با یک بافر از پیش تایپ‌شده که نوع بازگشتی آن از قبل مشخص است، به Match ارسال کرد. یک بافر از پیش تایپ‌شده نیز می‌تواند تأیید شود؛ این معمولاً برای بررسی روابط شکل مفید است:

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

نامی که به عنوان اولین آرگومان به Match and Verify ارسال می‌شود، الزامی است و به همراه توضیحی از محدودیت ناموفق، در خطاها نیز لحاظ می‌شود.

API خارجی FFI در xla/ffi/api/ffi.h مانند مثال‌های بالا، Error from Verify و ErrorOr<T> from Match را برمی‌گرداند. API داخلی FFI در xla/ffi/ffi.h همان رابط تطبیق را با استفاده از absl::Status و 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())));

آرگومان‌ها و نتایج متغیر

اگر تعداد آرگومان‌ها و نتیجه در نمونه‌های مختلف یک فراخوانی سفارشی متفاوت باشد، می‌توان آن‌ها را در زمان اجرا با استفاده از RemainingArgs و 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();
    });

آرگومان‌ها و نتایج متغیر می‌توانند پس از آرگومان‌ها و نتایج معمولی تعریف شوند، با این حال، اتصال آرگومان‌ها و نتایج معمولی پس از آرگومان‌ها و نتایج متغیر غیرقانونی است.

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

ویژگی‌ها

XLA FFI از رمزگشایی خودکار mlir::DictionaryAttr که به عنوان custom_call backend_config به آرگومان‌های کنترل‌کننده 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>

در این مثال، فراخوانی سفارشی یک آرگومان بافر و دو ویژگی دارد و XLA FFI می‌تواند به‌طور خودکار آنها را رمزگشایی کرده و به فراخوانی تعریف‌شده توسط کاربر منتقل کند.

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

ویژگی‌های Enum تعریف‌شده توسط کاربر

XLA FFI می‌تواند به طور خودکار ویژگی‌های انتگرالی MLIR را به enumهای تعریف شده توسط کاربر رمزگشایی کند. کلاس Enum باید نوع انتگرالی زیربنایی یکسانی داشته باشد و رمزگشایی باید به صراحت در 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(); });

اتصال تمام ویژگی‌های فراخوانی سفارشی

می‌توان به تمام ویژگی‌های فراخوانی سفارشی به عنوان یک دیکشنری دسترسی پیدا کرد و فقط ویژگی‌هایی را که در زمان اجرا مورد نیاز هستند، به صورت تنبل رمزگشایی کرد.

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

ویژگی‌های ساختار تعریف‌شده توسط کاربر

XLA FFI می‌تواند ویژگی‌های دیکشنری را به ساختارهای تعریف‌شده توسط کاربر رمزگشایی کند.

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

در مثال بالا، range یک ویژگی mlir::DictionaryAttr است و به جای دسترسی به فیلدهای دیکشنری از طریق نام، می‌تواند به طور خودکار به عنوان یک ساختار C++ رمزگشایی شود. رمزگشایی باید به طور صریح با یک ماکروی XLA_FFI_REGISTER_STRUCT_ATTR_DECODING ثبت شود (در پشت صحنه، یک قالب تخصصی در فضای نام ::xla::ffi تعریف می‌کند، بنابراین ماکرو باید به فضای نام سراسری اضافه شود).

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

ویژگی‌های سفارشی را می‌توان مانند هر ویژگی دیگری از یک دیکشنری بارگذاری کرد. در مثال زیر، تمام ویژگی‌های فراخوانی سفارشی به عنوان یک Dictionary رمزگشایی می‌شوند و یک range می‌توان با نام آن دسترسی داشت.

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

ایجاد یک فراخوانی سفارشی روی CPU

شما می‌توانید یک دستورالعمل HLO ایجاد کنید که نشان‌دهنده یک فراخوانی سفارشی از طریق API کلاینت XLA باشد. برای مثال، کد زیر از یک فراخوانی سفارشی برای محاسبه A[i] = B[i % 128]+ C[i] در CPU استفاده می‌کند. (البته شما می‌توانید - و باید! - این کار را با HLO معمولی انجام دهید.)

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

ایجاد یک فراخوانی سفارشی روی GPU

ثبت فراخوانی سفارشی GPU با XLA FFI تقریباً یکسان است، تنها تفاوت این است که برای GPU باید یک جریان پلتفرم زیربنایی (جریان CUDA یا ROCM) را درخواست کنید تا بتوانید هسته را روی دستگاه راه‌اندازی کنید. در اینجا یک مثال 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 = 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);

ابتدا توجه کنید که تابع فراخوانی سفارشی GPU هنوز تابعی است که روی CPU اجرا می‌شود . تابع do_custom_call CPU مسئول قرار دادن کارها در صف GPU است. در اینجا یک هسته CUDA را راه‌اندازی می‌کند، اما می‌تواند کار دیگری مانند فراخوانی cuBLAS را نیز انجام دهد.

آرگومان‌ها و نتایج نیز روی میزبان وجود دارند و عضو داده شامل یک اشاره‌گر به حافظه دستگاه (یعنی GPU) است. بافرهای ارسالی به کنترل‌کننده فراخوانی سفارشی، شکل بافرهای دستگاه اصلی را دارند، بنابراین فراخوانی سفارشی می‌تواند پارامترهای راه‌اندازی هسته را از آنها محاسبه کند.

ارسال تاپل‌ها به فراخوانی‌های سفارشی

فراخوانی سفارشی زیر را در نظر بگیرید.

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

در هر دو CPU و GPU، یک تاپل در حافظه به صورت آرایه‌ای از اشاره‌گرها نمایش داده می‌شود. وقتی XLA فراخوانی‌های سفارشی با آرگومان‌ها یا نتایج تاپل را فراخوانی می‌کند، آنها را مسطح کرده و به عنوان آرگومان‌ها یا نتایج بافر معمولی ارسال می‌کند.

خروجی‌های تاپل به عنوان بافرهای موقت

ورودی‌های تاپل به فراخوانی‌های سفارشی یک مزیت هستند، اما کاملاً ضروری نیستند. اگر ما از ورودی‌های تاپل به فراخوانی‌های سفارشی پشتیبانی نمی‌کردیم، همیشه می‌توانستید تاپل‌ها را با استفاده از get-tuple-element قبل از ارسال به فراخوانی سفارشی، از حالت فشرده خارج کنید.

از طرف دیگر، خروجی‌های چندتایی به شما امکان می‌دهند کارهایی را انجام دهید که در غیر این صورت نمی‌توانستید.

دلیل واضح داشتن خروجی‌های چندتایی این است که خروجی‌های چندتایی نحوه‌ی برگرداندن چندین آرایه‌ی مستقل توسط یک فراخوانی سفارشی (یا هر عملیات XLA دیگر) را نشان می‌دهند.

اما نکته‌ی کمتر واضح این است که یک خروجی چندتایی همچنین راهی برای دادن حافظه موقت به فراخوانی سفارشی شماست. بله، یک خروجی می‌تواند نشان‌دهنده‌ی یک بافر موقت باشد. در نظر بگیرید، یک بافر خروجی این ویژگی را دارد که اپراتور می‌تواند در آن بنویسد و پس از نوشتن در آن، می‌تواند از آن بخواند. این دقیقاً همان چیزی است که شما از یک بافر موقت می‌خواهید.

در مثال بالا، فرض کنید می‌خواهیم از F32[1024] به عنوان یک بافر موقت استفاده کنیم. در این صورت HLO را دقیقاً مانند بالا می‌نویسیم و هرگز اندیس تاپل ۱ خروجی فراخوانی سفارشی را نمی‌خوانیم.