Псевдонимы в XLA

В этом документе описывается 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 был передан во время выполнения).