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

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

Чтобы устранить первый недостаток, мы предлагаем ввести последний набор новых асинхронных опкодов: kAsyncStart , kAsyncUpdate и kAsyncDone . Идея состоит в том, чтобы создать общий асинхронный код операции, который может обернуть любую инструкцию HLO. Фактическая операция, которая будет выполняться асинхронно, будет закодирована с использованием вызываемого вычисления, корнем которого является только инструкция и любые входные параметры. Обработку и псевдонимы буфера ввода-вывода в реальном времени можно затем использовать совместно для любой асинхронной операции. Форма вывода инструкции асинхронного запуска тогда будет кортежем входных операндов, выходных значений и любого промежуточного состояния, необходимого для инструкций 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_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)

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

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

%op-start = (f32[64], f32[32], s32[]) op-start(f32[64] %operand),
                                      op_specific_attr=”foo”
%op-update0 = (f32[64], f32[32], s32[]) op-update(
                        (f32[64], f32[32], s32[]) %op-start),
                        op_specific_attr=”foo”
%op-update1 = (f32[64], f32[32], s32[]) op-update(
                        (f32[64], f32[32], s32[]) %op-update0)
%op-done = f32[32] op-done((f32[64], f32[32], s32[]) %op-update1)

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