XLA-Anfragen

In diesem Dokument wird beschrieben, wie Sie benutzerdefinierte XLA-Aufrufe mit der XLA FFI-Bibliothek schreiben und verwenden. Ein benutzerdefinierter Aufruf ist ein Mechanismus, um dem XLA-Compiler (zur Kompilierungszeit) einen externen „Vorgang“ im HLO-Modul zu beschreiben. XLA FFI ist ein Mechanismus, um die Implementierung solcher Vorgänge bei XLA (zur Laufzeit) zu registrieren. FFI steht für „Foreign Function Interface“ und ist eine Reihe von C-APIs, die eine binäre Schnittstelle (ABI) für XLA definieren, um externen Code aufzurufen, der in anderen Programmiersprachen geschrieben wurde. XLA bietet reine Header-Bindungen für XLA FFI, die in C++ geschrieben wurden und alle Details der zugrunde liegenden C-APIs für den Endnutzer ausblenden.

Benutzerdefinierte Aufrufe in JAX + XLA

In der JAX-Dokumentation finden Sie End-to-End-Beispiele für die Integration von benutzerdefinierten Aufrufen und XLA FFI in JAX.

XLA FFI-Bindung

Die XLA FFI-Bindung ist eine Spezifikation der benutzerdefinierten Aufrufsignatur zur Kompilierungszeit: benutzerdefinierte Aufrufargumente, Attribute und ihre Typen sowie zusätzliche Parameter, die über den Ausführungskontext übergeben werden (z.B. GPU-Stream für das GPU-Backend). Die XLA FFI-Bindung kann an jede aufrufbare C++-Funktion (Funktionszeiger, Lambda usw.) mit einer kompatiblen operator()-Signatur gebunden werden. Der erstellte Handler decodiert den XLA FFI-Aufrufrahmen (definiert durch die stabile C-API), überprüft alle Parameter auf den Typ und leitet die decodierten Ergebnisse an den benutzerdefinierten Callback weiter.

Die XLA FFI-Bindung basiert stark auf der Metaprogrammierung von Vorlagen, um den erstellten Handler in den effizientesten Maschinencode zu kompilieren. Die Laufzeit-Overheads liegen in der Größenordnung von einigen Nanosekunden für jeden benutzerdefinierten Aufrufparameter.

XLA FFI-Anpassungspunkte werden als Vorlagenspezialisierungen implementiert.Nutzer können definieren, wie ihre benutzerdefinierten Typen decodiert werden sollen. So ist es beispielsweise möglich, eine benutzerdefinierte Decodierung für benutzerdefinierte enum class-Typen zu definieren.

Fehler von benutzerdefinierten Aufrufen zurückgeben

Benutzerdefinierte Aufrufimplementierungen müssen den Wert xla::ffi::Error zurückgeben, um der XLA-Laufzeit Erfolg oder Fehler zu signalisieren. Er ähnelt absl::Status und hat dieselben Fehlercodes. Wir verwenden absl::Status nicht, da es keine stabile ABI hat und es unsicher wäre, es zwischen der dynamisch geladenen benutzerdefinierten Aufrufbibliothek und XLA selbst zu übergeben.

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

Pufferargumente und -ergebnisse

XLA verwendet den Stil der Zielübergabe für Ergebnisse: Benutzerdefinierte Aufrufe (oder andere XLA-Vorgänge) weisen keinen Speicher für Ergebnisse zu, sondern schreiben stattdessen in Ziele, die von der XLA-Laufzeit übergeben werden. XLA verwendet die statische Pufferzuweisung und weist Puffer für alle Werte basierend auf ihren Live-Bereichen zur Kompilierungszeit zu.

Ergebnisse, die an FFI-Handler übergeben werden, sind in eine Result<T> Vorlage eingebunden, die eine zeigerähnliche Semantik hat: operator-> ermöglicht den Zugriff auf den zugrunde liegenden Parameter.

AnyBuffer -Argumente und -Ergebnisse ermöglichen den Zugriff auf benutzerdefinierte Aufrufpufferparameter beliebigen Datentyps. Dies ist nützlich, wenn der benutzerdefinierte Aufruf eine generische Implementierung hat, die für mehrere Datentypen funktioniert, und die benutzerdefinierte Aufrufimplementierung die Laufzeitverteilung basierend auf dem Datentyp durchführt. AnyBuffer ermöglicht den Zugriff auf den Pufferdatentyp, die Dimensionen und einen Zeiger auf den Puffer selbst.

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

Eingeschränkte Pufferargumente und -ergebnisse

Mit Buffer können Einschränkungen für den Pufferdatentyp und die Anzahl der Dimensionen hinzugefügt werden. Diese werden automatisch vom Handler überprüft und geben einen Fehler an die XLA-Laufzeit zurück, wenn die Laufzeitargumente nicht mit der FFI-Handlersignatur übereinstimmen.

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

Puffer abgleichen und überprüfen

Puffermuster können Einschränkungen ausdrücken, die spezifischer sind als die Argument- und Ergebnistypen in einer FFI-Bindung. Sie sind nützlich, um einen AnyBuffer auf einen konkreten Puffertyp zu verfeinern, Dimensionsgrößen zu überprüfen und Beziehungen zwischen mehreren Pufferformen zu prüfen.

Muster sind unveränderlich und können den Datentyp und den Rang entweder direkt oder mit den entsprechenden Modifikatoren angeben:

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

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

Die Funktion Match überprüft ein Muster und gibt den am stärksten eingeschränkten Puffertyp zurück. Um einen AnyBuffer zu verfeinern, muss das Muster genau einen Datentyp und einen Rang angeben. Beim Abgleichen eines bereits typisierten Puffers werden alle zusätzlichen Einschränkungen überprüft und der Typ beibehalten.

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

In diesem Beispiel werden beim Abgleichen von input die beiden Dimensionsgrößen erfasst und beim Abgleichen von output wird überprüft, ob es dieselbe Form hat. Erfassungen werden nur übernommen, wenn das vollständige Muster erfolgreich abgeglichen wurde.

Ein Dimensionsmuster kann uneingeschränkt, fest oder erfasst sein. Feste Werte und Erfassungszeiger werden implizit in Dimensionsmuster konvertiert:

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 beschreibt eine vollständige Form und legt den Rang anhand der Anzahl der Argumente fest. WithDim<I> oder WithDim(index, ...) schränkt Dimensionen positionsbezogen ein, ohne den vollständigen Rang festzulegen. Eine positionsbezogene Einschränkung erfordert, dass die referenzierte Dimension vorhanden ist, lässt aber alle anderen Dimensionen uneingeschränkt:

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 erfordert einen Rang-4-F32-Puffer, dessen Dimensionen 0 und 3 gleich 1 bzw. 128 sind. Die Dimensionen 1 und 2 sind uneingeschränkt. Ohne ein explizites WithRank, akzeptiert WithDim(index, ...) jeden Rang, der größer als index ist.

Dimensionserfassungen können auch Beziehungen zwischen Positionen ausdrücken:

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

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

Wenn derselbe Erfassungszeiger mehrmals vorkommt, müssen alle entsprechenden Dimensionen dieselbe Größe haben. Der Aufruf Match oben akzeptiert daher nur quadratische Matrizen. Dimensionserfassungen werden erst geschrieben, nachdem das vollständige Muster erfolgreich abgeglichen wurde. Bei einem Fehler bleiben sie unverändert.

Für Beziehungen zwischen ganzen Puffern ist es nicht erforderlich, jede Dimension zu erfassen. WithShapeOf gleicht die vollständige Laufzeitform eines anderen Puffers ab, erlaubt aber einen anderen Datentyp. Like erfordert auch denselben Datentyp:

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

Beide Modifikatoren kopieren die Metadaten des Referenzpuffers, wenn das Muster erstellt wird. Das Muster behält keinen Verweis auf den Puffer bei. Da es sich um Laufzeiteinschränkungen handelt, wird ein AnyBuffer nicht auf einen konkreten Rückgabetyp verfeinert. Sie sind hauptsächlich mit Verify oder mit Match nützlich, wenn der Eingabepuffer bereits einen konkreten Typ hat.

Verwenden Sie Verify, wenn der Puffer seinen vorhandenen Typ beibehalten soll oder wenn ein Muster mehr als einen Datentyp oder Rang akzeptiert. Ein Indexpuffer kann beispielsweise entweder S32 oder S64 akzeptieren und nach der Überprüfung anhand des Datentyps verteilen:

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

Da dieses Muster mehrere konkrete Puffertypen zulässt, kann ein AnyBuffer nicht mit Match verfeinert werden: Es gibt keinen eindeutigen Rückgabetyp. Es kann trotzdem mit einem bereits typisierten Puffer an Match übergeben werden, dessen Rückgabetyp bereits bekannt ist. Ein bereits typisierter Puffer kann auch überprüft werden. Dies ist in der Regel nützlich, um Formbeziehungen zu prüfen:

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

Der Name, der als erstes Argument an Match und Verify übergeben wird, ist erforderlich und wird zusammen mit einer Beschreibung der fehlgeschlagenen Einschränkung in Fehler aufgenommen.

Die externe FFI-API in xla/ffi/api/ffi.h gibt Error von Verify und ErrorOr<T> von Match zurück, wie in den obigen Beispielen. Die interne FFI-API in xla/ffi/ffi.h bietet dieselbe Abgleichsschnittstelle mit absl::Status und 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())));

Variadische Argumente und Ergebnisse

Wenn sich die Anzahl der Argumente und Ergebnisse in verschiedenen Instanzen eines benutzerdefinierten Aufrufs unterscheiden kann, können sie zur Laufzeit mit RemainingArgs und RemainingRets decodiert werden.

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

Variadische Argumente und Ergebnisse können nach regulären Argumenten und Ergebnissen deklariert werden. Es ist jedoch nicht zulässig, reguläre Argumente und Ergebnisse nach variadischen zu binden.

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

Attribute

XLA FFI unterstützt die automatische Decodierung von mlir::DictionaryAttr, das als backend_config von custom_call übergeben wird, in FFI-Handlerargumente.

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

In diesem Beispiel hat der benutzerdefinierte Aufruf ein einzelnes Pufferargument und zwei Attribute. XLA FFI kann sie automatisch decodieren und an die benutzerdefinierte aufrufbare Funktion übergeben.

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

Benutzerdefinierte Enum-Attribute

XLA FFI kann integrale MLIR-Attribute automatisch in benutzerdefinierte Enums decodieren. Die Enum-Klasse muss denselben zugrunde liegenden integralen Typ haben und die Decodierung muss explizit bei XLA FFI registriert werden.

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

Alle benutzerdefinierten Aufrufattribute binden

Es ist möglich, als Dictionary auf alle benutzerdefinierten Aufrufattribute zuzugreifen und nur die Attribute verzögert zu decodieren, die zur Laufzeit benötigt werden.

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

Benutzerdefinierte Struct-Attribute

XLA FFI kann Dictionary-Attribute in benutzerdefinierte Structs decodieren.

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

Im obigen Beispiel ist range ein mlir::DictionaryAttr-Attribut. Anstatt über den Namen auf Dictionary-Felder zuzugreifen, kann es automatisch als C++-Struct decodiert werden. Die Decodierung muss explizit mit einem XLA_FFI_REGISTER_STRUCT_ATTR_DECODING-Makro registriert werden. Im Hintergrund wird eine Vorlagenspezialisierung im Namespace ::xla::ffi definiert. Daher muss das Makro dem globalen Namespace hinzugefügt werden.

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

Benutzerdefinierte Attribute können wie jedes andere Attribut aus einem Dictionary geladen werden. Im folgenden Beispiel werden alle benutzerdefinierten Aufrufattribute als Dictionary decodiert und auf range kann über den Namen zugegriffen werden.

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

Benutzerdefinierten Aufruf auf der CPU erstellen

Sie können mit der Client-API von XLA eine HLO-Anweisung erstellen, die einen benutzerdefinierten Aufruf darstellt. Im folgenden Code wird beispielsweise ein benutzerdefinierter Aufruf verwendet, um A[i] = B[i % 128]+ C[i] auf der CPU zu berechnen. Natürlich können und sollten Sie das mit regulärem HLO tun. –

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

Benutzerdefinierten Aufruf auf der GPU erstellen

Die Registrierung benutzerdefinierter Aufrufe auf der GPU mit XLA FFI ist fast identisch. Der einzige Unterschied besteht darin, dass Sie für die GPU einen zugrunde liegenden Plattformstream (CUDA- oder ROCM-Stream) anfordern müssen, um den Kernel auf dem Gerät zu starten. Hier ist ein CUDA-Beispiel, das dieselbe Berechnung (A[i] = B[i % 128] + C[i]) wie der CPU-Code oben durchführt.

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

Beachten Sie, dass die benutzerdefinierte Aufruffunktion für die GPU immer noch eine Funktion ist, die auf der CPU ausgeführt wird. Die CPU-Funktion do_custom_call ist für das Einreihen von Aufgaben auf der GPU verantwortlich. Hier wird ein CUDA-Kernel gestartet, aber es könnte auch etwas anderes ausgeführt werden, z. B. ein Aufruf von cuBLAS.

Argumente und Ergebnisse befinden sich ebenfalls auf dem Host und das Datenelement enthält einen Zeiger auf den Gerätespeicher (d.h. GPU). Die an den benutzerdefinierten Aufrufhandler übergebenen Puffer haben die Form der zugrunde liegenden Gerätepuffer, sodass der benutzerdefinierte Aufruf Kernel-Startparameter daraus berechnen kann.

Tupel an benutzerdefinierte Aufrufe übergeben

Betrachten Sie den folgenden benutzerdefinierten Aufruf.

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

Sowohl auf der CPU als auch auf der GPU wird ein Tupel im Speicher als Array von Zeigern dargestellt. Wenn XLA benutzerdefinierte Aufrufe mit Tupelargumenten oder -ergebnissen aufruft, werden diese vereinfacht und als reguläre Pufferargumente oder -ergebnisse übergeben.

Tupelausgaben als temporäre Puffer

Tupel als Eingaben für benutzerdefinierte Aufrufe sind praktisch, aber nicht unbedingt erforderlich. Wenn Tupel als Eingaben für benutzerdefinierte Aufrufe nicht unterstützt würden, könnten Sie die Tupel immer mit „get-tuple-element“ entpacken, bevor Sie sie an den benutzerdefinierten Aufruf übergeben.

Mit Tupel-Ausgaben können Sie jedoch Dinge tun, die sonst nicht möglich wären.

Der offensichtliche Grund für Tupelausgaben ist, dass ein benutzerdefinierter Aufruf (oder ein anderer XLA-Vorgang) damit mehrere unabhängige Arrays zurückgeben kann.

Weniger offensichtlich ist, dass eine Tupelausgabe auch eine Möglichkeit ist, dem benutzerdefinierten Aufruf temporären Speicher zur Verfügung zu stellen. Ja, eine Ausgabe kann einen temporären Puffer darstellen. Ein Ausgabepuffer hat die Eigenschaft, dass der Vorgang in ihn schreiben und daraus lesen kann, nachdem in ihn geschrieben wurde. Genau das ist es, was Sie von einem temporären Puffer erwarten.

Angenommen, wir möchten im obigen Beispiel F32[1024] als temporären Puffer verwenden. Dann schreiben wir das HLO wie oben und lesen einfach nie den Tupelindex 1 der Ausgabe des benutzerdefinierten Aufrufs.