XLA Custom Calls

Z tego dokumentu dowiesz się, jak pisać i używać niestandardowych wywołań XLA za pomocą biblioteki XLA FFI. Niestandardowe wywołanie to mechanizm opisujący zewnętrzną „operację” w module HLO dla kompilatora XLA (w czasie kompilacji), a XLA FFI to mechanizm rejestrowania implementacji takich operacji w XLA (w czasie wykonywania). FFI to skrót od „foreign function interface” (interfejs funkcji zewnętrznej). Jest to zestaw interfejsów API w języku C, które definiują interfejs binarny (ABI) umożliwiający XLA wywoływanie kodu zewnętrznego napisanego w innych językach programowania. XLA udostępnia powiązania tylko z nagłówkami dla XLA FFI napisane w C++, które ukrywają przed użytkownikiem wszystkie szczegóły niskiego poziomu dotyczące bazowych interfejsów API w języku C.

Niestandardowe wywołania JAX + XLA

Przykłady kompleksowej integracji wywołań niestandardowych i XLA FFI z JAX znajdziesz w dokumentacji JAX.

Powiązanie XLA FFI

Powiązanie XLA FFI to specyfikacja sygnatury wywołania niestandardowego w czasie kompilacji: argumenty wywołania niestandardowego, atrybuty i ich typy oraz dodatkowe parametry przekazywane przez kontekst wykonania (np. strumień GPU dla backendu GPU). Powiązanie XLA FFI można powiązać z dowolnym wywoływalnym elementem C++ (wskaźnikiem funkcji, lambdą itp.) z kompatybilną sygnaturą operator(). Skonstruowany moduł obsługi dekoduje ramkę wywołania XLA FFI (zdefiniowaną przez stabilny interfejs API w języku C), sprawdza typy wszystkich parametrów i przekazuje zdekodowane wyniki do zdefiniowanego przez użytkownika wywołania zwrotnego.

Powiązanie XLA FFI w dużym stopniu opiera się na metaprogramowaniu szablonów, aby umożliwić kompilowanie skonstruowanego modułu obsługi do najbardziej wydajnego kodu maszynowego. Obciążenie w czasie wykonywania wynosi kilka nanosekund na każdy parametr wywołania niestandardowego.

Punkty dostosowywania XLA FFI są implementowane jako specjalizacje szablonów, a użytkownicy mogą definiować sposób dekodowania typów niestandardowych. Można na przykład zdefiniować niestandardowe dekodowanie typów enum class zdefiniowanych przez użytkownika.

Zwracanie błędów z wywołań niestandardowych

Implementacje wywołań niestandardowych muszą zwracać wartość xla::ffi::Error, aby sygnalizować powodzenie lub błąd środowisku wykonawczemu XLA. Jest to podobne do absl::Status i ma ten sam zestaw kodów błędów. Nie używamy absl::Status, ponieważ nie ma stabilnego ABI i przekazywanie go między dynamicznie ładowaną biblioteką wywołań niestandardowych a samym XLA byłoby niebezpieczne.

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

Argumenty i wyniki buforów

XLA używa stylu przekazywania do miejsca docelowego w przypadku wyników: wywołania niestandardowe (lub inne operacje XLA) nie przydzielają pamięci na wyniki, ale zapisują je w miejscach docelowych przekazywanych przez środowisko wykonawcze XLA. XLA używa statycznego przypisywania buforów i przydziela bufory dla wszystkich wartości na podstawie ich zakresów aktywności w czasie kompilacji.

Wyniki przekazywane do modułów obsługi FFI są opakowane w szablon Result<T>, który ma semantykę podobną do wskaźnika: operator-> zapewnia dostęp do bazowego parametru.

Argumenty i wyniki AnyBuffer zapewniają dostęp do parametrów buforów wywołań niestandardowych dowolnego typu danych. Jest to przydatne, gdy wywołanie niestandardowe ma implementację ogólną, która działa w przypadku wielu typów danych, a implementacja wywołania niestandardowego wykonuje wysyłanie w czasie wykonywania na podstawie typu danych. AnyBuffer zapewnia dostęp do typu danych bufora, wymiarów i wskaźnika do samego bufora.

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

Argumenty i wyniki buforów z ograniczeniami

Buffer umożliwia dodawanie ograniczeń dotyczących typu danych bufora i liczby wymiarów. Moduł obsługi będzie je automatycznie sprawdzać i zwracać błąd do środowiska wykonawczego XLA, jeśli argumenty w czasie wykonywania nie będą zgodne z sygnaturą modułu obsługi 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();
    });

Dopasowywanie i weryfikowanie buforów

Wzorce buforów mogą wyrażać ograniczenia, które są bardziej szczegółowe niż typy argumentów i wyników w powiązaniu FFI. Są one przydatne do doprecyzowania AnyBuffer do konkretnego typu bufora, sprawdzania rozmiarów wymiarów i sprawdzania relacji między kształtami wielu buforów.

Wzorce są niezmienne i mogą określać typ danych i rangę bezpośrednio lub za pomocą odpowiednich modyfikatorów:

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

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

Funkcja Match sprawdza wzorzec i zwraca jego maksymalnie ograniczony typ bufora. Aby doprecyzować AnyBuffer, wzorzec musi określać dokładnie 1 typ danych i 1 rangę. Dopasowanie bufora, który ma już określony typ, weryfikuje wszelkie dodatkowe ograniczenia i zachowuje jego typ.

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

W tym przykładzie dopasowanie input przechwytuje rozmiary 2 wymiarów, a dopasowanie output sprawdza, czy ma on ten sam kształt. Przechwytywanie jest zatwierdzane tylko wtedy, gdy cały wzorzec zostanie dopasowany.

Wzorzec wymiaru może być nieograniczony, stały lub przechwycony. Wartości stałe i wskaźniki przechwytywania są niejawnie konwertowane na wzorce wymiarów:

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 opisuje pełny kształt i ustawia jego rangę na podstawie liczby argumentów. WithDim<I> lub WithDim(index, ...) ogranicza wymiary pozycyjnie bez ustalania pełnej rangi. Ograniczenie pozycyjne wymaga, aby istniał wymiar, do którego się odwołujesz, ale pozostawia wszystkie inne wymiary bez ograniczeń:

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 wymaga bufora F32 o randze 4, w którym wymiary 0 i 3 są równe odpowiednio 1 i 128, a wymiary 1 i 2 są nieograniczone. Bez jawnego WithRank, WithDim(index, ...) akceptuje dowolną rangę większą niż index.

Przechwytywanie wymiarów może też wyrażać relacje między pozycjami:

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

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

Gdy ten sam wskaźnik przechwytywania występuje wiele razy, wszystkie odpowiadające mu wymiary muszą mieć ten sam rozmiar. Wywołanie Match powyżej akceptuje więc tylko macierze kwadratowe. Przechwytywanie wymiarów jest zapisywane tylko wtedy, gdy cały wzorzec zostanie dopasowany, i nie zmienia się w przypadku niepowodzenia.

Relacje między całymi buforami nie wymagają przechwytywania każdego wymiaru. WithShapeOf dopasowuje pełny kształt w czasie wykonywania innego bufora, umożliwiając jednocześnie inny typ danych. Like wymaga też tego samego typu danych:

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

Oba modyfikatory kopiują metadane bufora odniesienia podczas tworzenia wzorca. Wzorzec nie zachowuje odniesienia do bufora. Ponieważ są to ograniczenia w czasie wykonywania, nie doprecyzowują one AnyBuffer do konkretnego typu zwracanego. Są one przydatne głównie w przypadku Verify lub Match, gdy bufor wejściowy ma już konkretny typ.

Użyj Verify, gdy bufor ma zachować swój dotychczasowy typ lub gdy wzorzec akceptuje więcej niż 1 typ danych lub rangę. Na przykład bufor indeksu może akceptować S32 lub S64 i wysyłać dane na podstawie typu danych po weryfikacji:

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

Ponieważ ten wzorzec dopuszcza wiele konkretnych typów buforów, nie może doprecyzować AnyBuffer za pomocą Match: nie ma unikalnego zwracanego typu. Można go jednak przekazać do Match z buforem, który ma już określony typ, którego zwracany typ jest już znany. Można też zweryfikować bufor, który ma już określony typ. Jest to zwykle przydatne do sprawdzania relacji kształtów:

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

Nazwa przekazywana jako pierwszy argument do Match i Verify jest wymagana i jest uwzględniana w błędach wraz z opisem nieudanego ograniczenia.

Zewnętrzny interfejs API FFI w xla/ffi/api/ffi.h zwraca Error z Verify i ErrorOr<T> z Match, jak w przykładach powyżej. Wewnętrzny interfejs API FFI w xla/ffi/ffi.h udostępnia ten sam interfejs dopasowywania za pomocą absl::Status i 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())));

Argumenty i wyniki zmienne

Jeśli liczba argumentów i wyników może się różnić w różnych instancjach wywołania niestandardowego, można je dekodować w czasie wykonywania za pomocą RemainingArgs i 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();
    });

Argumenty i wyniki zmienne można deklarować po zwykłych argumentach i wynikach, ale powiązanie zwykłych argumentów i wyników po argumentach i wynikach zmiennych jest niedozwolone.

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

Atrybuty

XLA FFI obsługuje automatyczne dekodowanie mlir::DictionaryAttr przekazywanego jako custom_call backend_config do argumentów modułu obsługi 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>

W tym przykładzie wywołanie niestandardowe ma 1 argument bufora i 2 atrybuty, a XLA FFI może je automatycznie dekodować i przekazywać do wywoływalnego elementu zdefiniowanego przez użytkownika.

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

Atrybuty wyliczenia zdefiniowane przez użytkownika

XLA FFI może automatycznie dekodować atrybuty całkowite MLIR na wyliczenia zdefiniowane przez użytkownika. Klasa wyliczenia musi mieć ten sam bazowy typ całkowity, a dekodowanie musi być jawnie zarejestrowane w 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(); });

Powiązanie wszystkich atrybutów wywołania niestandardowego

Można uzyskać dostęp do wszystkich atrybutów wywołania niestandardowego jako słownika i leniwie dekodować tylko te atrybuty, które są potrzebne w czasie wykonywania.

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

Atrybuty struktury zdefiniowane przez użytkownika

XLA FFI może dekodować atrybuty słownika na struktury zdefiniowane przez użytkownika.

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

W przykładzie powyżej range to atrybut mlir::DictionaryAttr. Zamiast uzyskiwać dostęp do pól słownika według nazwy, można go automatycznie dekodować jako strukturę C++. Dekodowanie musi być jawnie zarejestrowane za pomocą makra XLA_FFI_REGISTER_STRUCT_ATTR_DECODING (w tle definiuje specjalizację szablonu w przestrzeni nazw ::xla::ffi, dlatego makro musi zostać dodane do globalnej przestrzeni nazw).

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

Atrybuty niestandardowe można wczytywać ze słownika tak jak każdy inny atrybut. W przykładzie poniżej wszystkie atrybuty wywołania niestandardowego są dekodowane jako Dictionary, a do range można uzyskać dostęp według nazwy.

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

Tworzenie wywołania niestandardowego na procesorze

Za pomocą interfejsu API klienta XLA możesz utworzyć instrukcję HLO, która reprezentuje wywołanie niestandardowe. Na przykład ten kod używa wywołania niestandardowego do obliczania A[i] = B[i % 128]+ C[i] na procesorze. (Oczywiście możesz – i powinieneś! – zrobić to za pomocą zwykłego 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);

Tworzenie wywołania niestandardowego na GPU

Rejestracja wywołania niestandardowego GPU w XLA FFI jest prawie identyczna. Jedyna różnica polega na tym, że w przypadku GPU musisz poprosić o bazowy strumień platformy (CUDA lub ROCM), aby móc uruchomić jądro na urządzeniu. Oto przykład CUDA, który wykonuje te same obliczenia (A[i] = B[i % 128] + C[i]) co kod procesora powyżej.

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

Zwróć uwagę, że funkcja wywołania niestandardowego GPU jest nadal funkcją wykonywaną na procesorze. Funkcja procesora do_custom_call odpowiada za umieszczanie pracy w kolejce na GPU. W tym przypadku uruchamia jądro CUDA, ale może też robić coś innego, np. wywoływać cuBLAS.

Argumenty i wyniki też znajdują się na hoście, a element danych zawiera wskaźnik do pamięci urządzenia (np. GPU). Bufory przekazywane do modułu obsługi wywołania niestandardowego mają kształt bazowych buforów urządzenia, dzięki czemu wywołanie niestandardowe może obliczać na ich podstawie parametry uruchamiania jądra.

Przekazywanie krotek do wywołań niestandardowych

Rozważ to wywołanie niestandardowe.

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

Zarówno na procesorze, jak i na GPU krotka jest reprezentowana w pamięci jako tablica wskaźników. Gdy XLA wywołuje wywołania niestandardowe z argumentami lub wynikami krotki, spłaszcza je i przekazuje jako zwykłe argumenty lub wyniki buforów.

Wyniki krotek jako bufory tymczasowe

Dane wejściowe krotek do wywołań niestandardowych są wygodne, ale nie są bezwzględnie konieczne. Gdybyśmy nie obsługiwali danych wejściowych krotek do wywołań niestandardowych, zawsze można by rozpakować krotki za pomocą get-tuple-element przed przekazaniem ich do wywołania niestandardowego.

Z drugiej strony wyniki krotek pozwalają robić rzeczy, których inaczej nie można by zrobić.

Oczywistym powodem, dla którego warto mieć wyniki krotek, jest to, że w ten sposób wywołanie niestandardowe (lub dowolna inna operacja XLA) zwraca wiele niezależnych tablic.

Mniej oczywiste jest jednak to, że wynik krotki to też sposób na udostępnienie wywołaniu niestandardowemu pamięci tymczasowej. Tak, wynik może reprezentować bufor tymczasowy. Rozważmy, że bufor wyjściowy ma tę właściwość, że operacja może do niego zapisywać i odczytywać z niego po zapisaniu. To dokładnie to, czego oczekujesz od bufora tymczasowego.

W przykładzie powyżej załóżmy, że chcemy użyć F32[1024] jako bufora tymczasowego. Wtedy napiszemy HLO tak jak powyżej i po prostu nigdy nie odczytamy indeksu krotki 1 z wyniku wywołania niestandardowego.