Асинхронные инструкции HLO

1. Добавление асинхронных операций в HLO — сложная задача (например, all-reduce-start и all-reduce-done ).
2. Разделение на этапы «начало» и «завершение» может оказаться недостаточным для некоторых сценариев асинхронного выполнения.

Для устранения первого недостатка мы предлагаем ввести еще один набор новых асинхронных кодов операций: kAsyncStart , kAsyncUpdate и kAsyncDone . Идея состоит в создании универсального асинхронного кода операции, который может обернуть любую инструкцию HLO. Фактическая операция, которая будет выполняться асинхронно, будет кодироваться с помощью вызываемого вычисления, которое имеет в качестве корня только инструкцию и любые параметры для входных данных. Обработка входных/выходных буферов и псевдонимы в процессе выполнения могут быть общими для любой асинхронной операции. Форма выходных данных инструкции async-start будет представлять собой кортеж из входных операндов, выходных значений и любого промежуточного состояния, необходимого для инструкций async-update или async-done .

%async_op {
  %param0 = f32[64] parameter(0)
  ROOT %op = f32[32] op(f32[64] %param0), op_specific_attr=”foo”
}

%async-start = ((f32[64]), f32[32], s32[]) async-start(f32[64] %operand),
                                           calls=%async_op
%async-done = f32[32] async-done(((f32[64]), f32[32], s32[]) %async-start)

В приведенном выше представлении только async-start содержит вызываемое вычисление, поскольку легко определить, что делает async-done проследив за его операндом, чтобы найти соответствующий async-start и, следовательно, вызываемое вычисление.

Также обратите внимание, что первый элемент в выходном кортеже async-start представляет собой кортеж, содержащий операнды. Элементы этого кортежа операндов имеют псевдонимы с соответствующими операндами, поэтому их буферы остаются активными как минимум до выполнения инструкции async-done . Аналогично, второй элемент имеет псевдоним с выходными данными инструкции async-done , а третий элемент представляет собой состояние контекста, используемое для отслеживания асинхронной операции. Это представление естественным образом поддерживает несколько тензоров во входных и/или выходных данных асинхронной операции:

%async_op {
  %param0 = f32[64] parameter(0)
  %param1 = f32[64] parameter(1)
  ROOT %op = (f32[32], f32[32]) op(f32[64] %param0, f32[64] %param1),
                                op_specific_attr=”foo”
}

%async-start = ((f32[64], f32[64]), (f32[32], f32[32]), s32[])
               async-start(f32[64] %operand0, f32[64] %operand1),
               calls=%async_op
%async-done = (f32[32], f32[32]) async-done(%async-start)

Кроме того, операция может быть дополнительно разложена на ноль или более шагов async-update , выполняющих промежуточные вычисления. Псевдонимы ввода/вывода работают аналогично инструкции async-update , и каждая инструкция async-start и async-update должна иметь одного пользователя, который является либо другой инструкцией async-update , либо async-done .

%async_op {
  %param0 = f32[64] parameter(0)
  ROOT %op = f32[32] op(f32[64] %param0), op_specific_attr=”foo”
}

%async-start = ((f32[64]), f32[32], s32[]) async-start(f32[64] %operand),
                                         calls=%async_op
%async-update0 = ((f32[64]), f32[32], s32[]) async-update(
                           ((f32[64]), f32[32], s32[]) %async-start)
%async-update1 = ((f32[64]), f32[32], s32[]) async-update(
                           ((f32[64]), f32[32], s32[]) %async-update0)
%async-done = f32[32] async-done(((f32[64]), f32[32], s32[]) %async-update1)

Синтаксический сахар

Парсер HLO поддерживает синтаксический сахар для автоматического анализа и вывода асинхронных операций, как если бы они были первоклассными кодами операций. Парсер обрабатывает суффиксы -start , -update и -done особым образом, автоматически создавая асинхронное вычисление и обернутую инструкцию (без суффикса).

Например, асинхронный custom-call можно записать следующим образом:

%cc-start = ((f32[64]), f32[32], s32[]) custom-call-start(%operand),
                                        custom_call_target="foo"
%cc-update = ((f32[64]), f32[32], s32[]) custom-call-update(%cc-start)
%result = f32[32] custom-call-done(%cc-update)

Парсер преобразует это в следующий эквивалентный HLO:

%async_computation {
  %p0 = f32[64] parameter(0)
  ROOT %custom-call = f32[32] custom-call(%p0), custom_call_target="foo"
}

%async-start = ((f32[64]), f32[32], s32[]) async-start(%operand),
                calls=%async_computation
%async-update = ((f32[64]), f32[32], s32[]) async-update(%async-start)
%result = f32[32] async-done(%async-update)

Эта функция десахаризации поддерживается для большинства опкодов HLO (например, custom-call , dot , all-reduce и т. д.).

Исключения

Во избежание неоднозначностей парсер не будет десахаризировать операции, имеющие явно определенные первоклассные опкоды с суффиксами -start и/или -done (например, copy-start / copy-done , collective-permute-start / collective-permute-done ). В этих операциях будут по-прежнему использоваться их соответствующие первоклассные опкоды.

Позднее связывание

В некоторых случаях операнды (входы) или выходы асинхронной операции не все доступны или выделены на момент начала операции. XLA поддерживает позднее связывание , которое позволяет связывать операнды постепенно во время шагов async-update , а выходы — во время шагов async-update или async-done .

Представительство в HLO

Для вызываемого вычисления, которое ожидает \(N\) С помощью этих параметров мы можем начать асинхронное выполнение, используя меньшее количество параметров, чем... \(N\) операнды. Остальные операнды передаются в последующих инструкциях async-update .

• async-start привязывает первый \(K\) операнды (\(K < N\)).
• Инструкции async-update связывают оставшиеся \(N - K\) операнды.

Привязка операндов должна происходить слева направо. То есть, если вычисление ожидает параметры. \(P_0, P_1, \dots, P_{N-1}\)Они должны быть связаны именно в таком порядке по всей асинхронной цепочке.

Формы async-start и async-update отражают постепенное связывание параметров. В частности, первый элемент кортежа (формы операндов) увеличивается по мере связывания большего количества операндов.

Привязка выходных данных не зависит от привязки операндов и может происходить на любом этапе асинхронной цепочки (либо при async-update , либо при завершении async-done ).

Пример с использованием kCall

Рассмотрим вычисляемую функцию %foo , которая принимает два параметра:

%foo {
  %p0 = f32[] parameter(0)
  %p1 = f32[] parameter(1)
  ROOT %add = f32[] add(%p0, %p1)
}

Мы можем вызвать эти вычисления асинхронно, привязав %p0 в начале и %p1 в момент обновления:

%call-start = ((f32[]), (), s32[]) call-start(%operand0), to_apply=%foo
%call-update = ((f32[], f32[]), f32[], s32[]) call-update(%call-start, %operand1)
%result = f32[] call-done(%call-update)

Парсер преобразует это в следующий HLO:

%async-start = ((f32[]), (), s32[]) async-start(%operand0), calls=%foo
%async-update = ((f32[], f32[]), f32[], s32[]) async-update(%async-start, %operand1)
%result = f32[] async-done(%async-update)

Выходы с поздней границей

Помимо операндов (входов), выходные данные асинхронной операции также могут быть привязаны с задержкой. Это полезно, когда выходные буферы неизвестны или выделены в начале операции.

Для представления результатов с поздним ограничением:

1. Инструкция async-start (или call-start ) определяется с помощью пустого кортежа () по индексу 1 в структуре выходных данных (слоте результата).
2. Последующая инструкция async-update (или call-update ) указывает фактическую форму выходных данных по индексу 1, заменяя пустой кортеж.
3. В качестве альтернативы, выходные данные могут быть привязаны в конце цепочки с помощью инструкции async-done (или call-done ), которая возвращает окончательную форму выходных данных. Это можно сделать независимо от наличия промежуточных шагов async-update в цепочке.

Пример с async-update

// Output is not bound at start (index 1 is ())
%call-start = ((f32[1024]), (), s32[]) call-start(%input_buffer), to_apply=%foo

// Output is bound at update (index 1 becomes (f32[1024]))
%call-update = ((f32[1024]), (f32[1024]), s32[]) call-update(%call-start, %output_buffer)

%result = (f32[1024]) call-done(%call-update)

Парсер преобразует это в:

%async-start = ((f32[1024]), (), s32[]) async-start(%input_buffer), calls=%foo
%async-update = ((f32[1024]), (f32[1024]), s32[]) async-update(%async-start, %output_buffer)
%result = (f32[1024]) async-done(%async-update)

Пример с async-done (без async-update )

Если промежуточные этапы обновления отсутствуют, вывод можно напрямую привязать в async-done :

// Output is not bound at start (index 1 is ())
%call-start = ((f32[1024]), (), s32[]) call-start(%input_buffer), to_apply=%foo

// Output is bound at done
%result = (f32[1024]) call-done(%call-start)

Парсер преобразует это в:

%async-start = ((f32[1024]), (), s32[]) async-start(%input_buffer), calls=%foo
%result = (f32[1024]) async-done(%async-start)

Пример с промежуточным async-update и привязкой вывода к моменту async-done

Если существуют промежуточные этапы обновления для привязки операндов, но выходной сигнал всё равно привязывается в самом конце:

// Output is not bound at start, no operands bound
%call-start = ((), (), s32[]) call-start(), to_apply=%foo

// Operands are bound at update, but output remains unbound (index 1 is ())
%call-update = ((f32[], f32[]), (), s32[]) call-update(%call-start, %operand0, %operand1)

// Output is bound at done
%result = f32[] call-done(%call-update)

Парсер преобразует это в:

%async-start = ((), (), s32[]) async-start(), calls=%foo
%async-update = ((f32[], f32[]), (), s32[]) async-update(%async-start, %operand0, %operand1)
%result = f32[] async-done(%async-update)