异步 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_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”操作码。