- HLO に非同期オペレーションを追加するのは面倒(
all-reduce-start
やall-reduce-done
)。 - 一部の非同期使用では、開始分割と完了分割が不適切である場合がある 対応できます
1 つ目の欠点に対処するため、最後にもう一つ、
非同期オペコード: 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 命令になるまでバッファは有効なままです。
同様に、2 番目の要素は async-done
の出力でエイリアスとなり、
3 つ目の要素はコンテキスト状態であり、コンテキストの
非同期オペレーションを実行しますこの表現では、複数のテンソルを
非同期オペレーションの入力/出力とエイリアス設定の動作は同じ
方法:
%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)
さらに、演算をさらに 0 個以上の 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)
構文糖
これから行われる演算を定義するために やや面倒です。Google では、Python でどのように 非同期オペレーションをファースト クラス オペレーションであるかのように自動的に出力および解析する 使用されます。接尾辞「-start」、「-update」、「-done」は、 計算と命令を自動的に作成することで(特別な処理を行わずに)、 使用します。たとえば、上記のコード スニペットは 2 つを解析して同じ表現に変換できます。
%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 用のオペコードを明示的に定義した場合は、async-start でラップする
「-start」や「-done」という接尾辞を付けます。これも
HLO レベルの処理を必要とする指示があった場合に、ハッチします。
上記のモデルに当てはまらない(例: エイリアス入出力
。したがって、最初に copy-start
/copy-done
が、
collective-permute-start
/collective-permute-done
などで引き続き使用
対応するファースト クラスのオペコードを
async-start
/async-done
個のオペコード(コードをクリーンアップして削除するまで)
「-start」/「-done」オペコード。