- Добавление асинхронных операций в HLO является громоздким (т.е.
all-reduce-start
иall-reduce-done
). - Разделение 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».
- Добавление асинхронных операций в HLO является громоздким (т.е.
all-reduce-start
иall-reduce-done
). - Разделение 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».