Aliasing in XLA

In diesem Dokument wird die XLA Aliasing API beschrieben, mit der Sie die Aliasing zwischen den Eingabe- und Ausgabepuffern beim Erstellen eines XLA-Programms.

Aliasing bei der Kompilierung definieren

Nehmen wir als Beispiel ein triviales HLO-Modul, das seiner Eingabe einfach 1 hinzufügt:

HloModule increment

ENTRY entry {
  %p = f32[] parameter(0)
  %c = f32[] constant(1)
  ROOT %out = f32[] add(%p, %c)
}

Dieses Modul weist zwei 4-Byte-Zwischenspeicher zu: einen für die Eingabe-%p und einen für die Eingabe. für die Ausgabe %out.

Häufig ist es jedoch sinnvoll, die Aktualisierung direkt durchzuführen, z. B. Im Frontend, das den Ausdruck generiert, ist die Eingabevariable nicht mehr aktiv nach der Berechnung, wie im Inkrement p++).

Um eine solche Aktualisierung effizient durchzuführen, können Sie das Eingabealiasing angeben:

HloModule increment, input_output_alias={ {}: 0 }

ENTRY entry {
  %p = f32[] parameter(0)
  %c = f32[] constant(1)
  ROOT %out = f32[] add(%p, %c)
}

Das Format gibt an, dass die gesamte Ausgabe (mit {} markiert) dem Alias Eingabeparameter 0.

Informationen zum programmatischen Festlegen des Aliasing finden Sie in der XlaBuilder::SetUpAlias der API erstellen.

Aliasing zur Laufzeit definieren

Das im vorherigen Schritt definierte Aliasing wird während der Kompilierung angegeben. Während der Ausführung können Sie den LocalClient::RunAsync API zur Auswahl, ob der Zwischenspeicher gespendet werden soll.

Eingabepuffer für das Programm sind in ExecutionInputs, die wiederum einen Baum von MaybeOwningDeviceMemory enthalten. Wenn der Arbeitsspeicher als owning angegeben (die Inhaberschaft des Zwischenspeichers wird an die XLA-Laufzeit übergeben), wird der Zwischenspeicher tatsächlich zur Verfügung gestellt und die Aktualisierung wird wie folgt ausgeführt: von der Compile-Time Aliasing API angefordert.

Wenn der Zwischenspeicher, der bei der Kompilierung mit einem Alias versehen wird, jedoch nicht am Laufzeit wird der Kopierschutz aktiviert: Ein zusätzlicher Ausgabepuffer O wird zugewiesen. und der Inhalt des Eingabepuffers P, der mit einem Alias versehen werden sollte, wird kopiert. in O (das Programm kann so effektiv ausgeführt werden, als ob der Zwischenspeicher O zur Laufzeit gespendet.)