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 ExecutionInput
s 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).