W tym dokumencie opisujemy interfejs XLA aliasing API, który pozwala określić aliasing między buforami danych wejściowych i wyjściowych podczas tworzenia programu XLA.
Definiowanie aliasu podczas kompilacji
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 bufory 4-bajtowe: jeden na dane wejściowe %p i jeden na dane wyjściowe %out.
Często jednak zalecane jest 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 p++).
Aby przeprowadzić taką aktualizację sprawnie, możesz określić aliasy wejściowe:
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 całe dane wyjściowe (oznaczone tagiem {}) są przypisywane aliasom parametru wejściowego 0.
Aby automatycznie określić aliasy, zapoznaj się z interfejsem API XlaBuilder::SetUpAlias.
Definiowanie aliasów w czasie działania
Alias zdefiniowany w poprzednim kroku jest określany podczas kompilacji.
Podczas wykonywania możesz za pomocą interfejsu API LocalClient::RunAsync zdecydować, czy przekazać bufor.
Bufory wejściowe programu są zawarte w komponencie ExecutionInput, który z kolei zawiera drzewo MaybeOwningDeviceMemory. Jeśli pamięć jest określona jako posiadająca pamięć (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 w czasie kompilacji.
Jeśli jednak bufor objęty aliasem podczas kompilacji nie jest przekazywany 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 pliku O (dzięki czemu program może działać tak, jakby bufor O został przekazany w czasie działania).