Aliasing in XLA

In diesem Dokument wird die XLA Aliasing API beschrieben, mit der Sie beim Erstellen eines XLA-Programms das Aliasing zwischen dem Eingabe- und Ausgabepuffer angeben können.

Aliasing bei der Kompilierung definieren

Betrachten Sie beispielsweise ein triviales HLO-Modul, das einfach 1 zur Eingabe 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-Puffer zu: einen für die Eingabe-%p und einen für die Ausgabe-%out.

Häufig ist es jedoch wünschenswert, die Aktualisierung direkt durchzuführen (z. B. wenn im Front-End, das den Ausdruck generiert, die Eingabevariable nach der Berechnung nicht mehr aktiv ist, wie in der Erhöhung von p++).

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

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 (gekennzeichnet mit {}) dem Eingabeparameter 0 zugeordnet wird.

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

Aliasing zur Laufzeit definieren

Das im vorherigen Schritt definierte Aliasing wird während der Kompilierung angegeben. Bei der Ausführung können Sie mit der LocalClient::RunAsync API auswählen, ob der Zwischenspeicher bereitgestellt werden soll.

Eingabepuffer für das Programm sind in ExecutionInputs eingeschlossen, die wiederum eine Baumstruktur von MaybeOwningDeviceMemory enthalten. Wenn für den Arbeitsspeicher owning festgelegt ist (die Eigentümerschaft des Zwischenspeichers wird an die XLA-Laufzeit übergeben), wird der Zwischenspeicher tatsächlich zur Verfügung gestellt und die Aktualisierung erfolgt, wie von der Compile-Time Aliasing API angefordert.

Wenn jedoch der Zwischenspeicher, der zum Kompilierungszeitpunkt mit Alias versehen wird, zur Laufzeit nicht zur Verfügung gestellt wird, wird copy-protection aktiviert: Es wird ein zusätzlicher Ausgabepuffer O zugewiesen und der Inhalt des Eingabepuffers P, der als Alias verwendet werden sollte, wird in O kopiert (so effektiv kann das Programm so ausgeführt werden, als wäre der Zwischenspeicher O zur Laufzeit gespendet).