В этом документе описывается API псевдонимов XLA, который позволяет указать псевдонимы между входным и выходным буферами при создании программы XLA.
Определение псевдонимов во время компиляции
Например, рассмотрим тривиальный модуль HLO, который просто добавляет 1
к своим входным данным:
HloModule increment
ENTRY entry {
%p = f32[] parameter(0)
%c = f32[] constant(1)
ROOT %out = f32[] add(%p, %c)
}
Этот модуль выделит два 4-байтовых буфера: один для ввода %p
и один для вывода %out
.
Однако часто желательно выполнить обновление на месте (например, если во внешнем интерфейсе, генерирующем выражение, входная переменная больше не активна после вычисления, как в случае приращения p++
).
Чтобы эффективно выполнить такое обновление, вы можете указать псевдоним ввода:
HloModule increment, input_output_alias={ {}: 0 }
ENTRY entry {
%p = f32[] parameter(0)
%c = f32[] constant(1)
ROOT %out = f32[] add(%p, %c)
}
Формат указывает, что весь вывод (отмеченный {}
) имеет псевдоним входного параметра 0
.
Чтобы программно указать псевдонимы, см. API XlaBuilder::SetUpAlias
.
Определение псевдонимов во время выполнения
Псевдонимы, определенные на предыдущем шаге, указываются во время компиляции . Во время выполнения вы можете использовать API LocalClient::RunAsync
чтобы выбрать, передавать ли буфер.
Входные буферы программы заключены в ExecutionInput
, которые, в свою очередь, содержат дерево MaybeOwningDeviceMemory
. Если память указана как владеющая (владение буфером передается среде выполнения XLA), буфер фактически передается в дар, и обновление выполняется на месте, как того требует API псевдонимов во время компиляции.
Если, однако, буфер, которому присвоен псевдоним во время компиляции, не передается во время выполнения, срабатывает защита от копирования : выделяется дополнительный выходной буфер O
, и содержимое входного буфера P
, которому должно было быть назначено псевдоним, копируется в O
(настолько эффективно, что программа может выполняться так, как если бы буфер O
был передан во время выполнения).