W tym dokumencie opisano interfejs API aliasingu XLA, który umożliwia określenie aliasów pomiędzy buforami danych wejściowych i wyjściowych podczas tworzenia programu XLA.
Definiowanie aliasu podczas kompilowania
Rozważmy na przykład prosty moduł HLO, który po prostu dodaje do danych wejściowych „1
”:
HloModule increment
ENTRY entry {
%p = f32[] parameter(0)
%c = f32[] constant(1)
ROOT %out = f32[] add(%p, %c)
}
Ten moduł przydzieli 2 4-bajtowe bufory: jeden na dane wejściowe %p
i jeden na dane wyjściowe %out
.
Często jednak zaleca się przeprowadzenie aktualizacji na miejscu (np. jeśli w interfejsie generującym wyrażenie zmienna wejściowa nie jest już aktywna po przeprowadzeniu obliczeń, jak w przypadku przyrostu wartości p++
).
Aby przeprowadzić taką aktualizację skutecznie, możesz określić aliasy danych wejściowych:
HloModule increment, input_output_alias={ {}: 0 }
ENTRY entry {
%p = f32[] parameter(0)
%c = f32[] constant(1)
ROOT %out = f32[] add(%p, %c)
}
Ten format określa, że wszystkie dane wyjściowe (oznaczone tagiem {}
) są przypisywane aliasem parametru wejściowego 0
.
Aby określić aliasy automatycznie, zapoznaj się z interfejsem API XlaBuilder::SetUpAlias
.
Definiowanie aliasów w czasie działania
Aliasy widoczne w poprzednim kroku są określane podczas kompilacji.
Podczas wykonywania możesz za pomocą interfejsu API LocalClient::RunAsync
zdecydować, czy przekazać bufor.
Bufory danych wejściowych programu są ujęte w pliki ExecutionInput
, które z kolei zawierają drzewo MaybeOwningDeviceMemory
. Jeśli pamięć jest określona jako posiadanie (własność bufora jest przekazywana do środowiska wykonawczego XLA), bufor jest faktycznie przekazywany, a aktualizacja jest wykonywana zgodnie z żądaniami interfejsu API do tworzenia aliasów podczas kompilacji.
Jeśli jednak bufor z aliasem podczas kompilacji nie jest udostępniany w czasie działania, uruchamia się ochrona przed kopiowaniem: przydzielony jest dodatkowy bufor wyjściowy O
, a zawartość bufora wejściowego P
, który miał być aliasem, jest kopiowana do usługi O
(dzięki czemu program może działać tak, jakby bufor O
został przekazany w czasie działania).