Aliasing w XLA

Ten dokument opisuje interfejs API do generowania aliasów XLA, który umożliwia określenie utworzenie aliasu między buforami wejściowymi i wyjściowymi podczas tworzenia programu XLA.

Definiowanie aliasów podczas kompilacji

Weźmy na przykład banalny 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 dwa 4-bajtowe bufory: jeden dla wejściowego %p i jeden dla wejścia dla danych wyjściowych %out.

Często przydaje się jednak przeprowadzanie aktualizacji na miejscu (np. we frontendzie podczas generowania wyrażenia zmienna wejściowa przestaje być aktywna po obliczeniu, jak w przyrostie 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)
}

Format określa, że całe dane wyjściowe (oznaczone przez {}) są powiązane aliasem z parametrem parametr wejściowy 0.

Aby automatycznie określić aliasy, zapoznaj się z dokumentacją XlaBuilder::SetUpAlias API.

Definiowanie aliasów w czasie działania

Aliasy zdefiniowane w poprzednim kroku są określane podczas kompilacji. Podczas wykonywania możesz użyć atrybutu LocalClient::RunAsync interfejsu API, aby określić, czy udostępnić bufor.

Bufory wejściowe programu są opakowane ExecutionInput, które z kolei zawierają drzewo MaybeOwningDeviceMemory. Jeśli pamięć to określona jako właściciel (własność bufora jest przekazywana do środowiska wykonawczego XLA), bufor jest faktycznie przekazywany, a aktualizacja jest przeprowadzana, żądania wysyłane przez interfejs API do generowania aliasów podczas kompilacji.

Jeśli jednak bufor wskazany jako alias podczas kompilacji nie zostanie przekazany na w środowisku wykonawczym, uruchamia się zabezpieczenie przed kopiowaniem: przydzielony jest dodatkowy bufor wyjściowy O, i zawartość bufora wejściowego P, który miał być aliasem, są kopiowane do O (tak skutecznie program może działać, jak gdyby bufor O był przekazywane w czasie działania).