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