非同步 HLO 操作說明

  1. 在 HLO 中新增非同步作業相當麻煩 (即 all-reduce-start 和 all-reduce-done)。
  2. 啟動和完成分割可能不適用於某些非同步使用 用途

為了鎖定第一次的短缺目標，建議貴公司介紹最後一組新的 非同步運算碼：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)

語法糖

由於透過獨立的運算定義 非同步執行中可能不太容易 自動列印及剖析非同步作業 Opcodes設計原則是將「-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)

為避免造成混淆，驗證器會禁止作業 如果我們明確為該程式碼 加上「-start」和/或「-done」後置字串之運算這也是跳脫 我們有任何需要 HLO 層級治療的指示， 不適用於上述模型 (例如別名的輸入/輸出) 緩衝區)。因此，一開始的 copy-start/copy-done collective-permute-start/collective-permute-done 等將繼續使用 各自獲得一流的營運代碼，而不是新的 async-start/async-done 運算程式碼，直到我們清理程式碼以移除這些程式碼 「-start」/"-done" 運算碼。