非同期 HLO の手順

  1. HLO に非同期オペレーションを追加するのは面倒(all-reduce-startall-reduce-done)。
  2. 一部の非同期使用では、開始分割と完了分割が不適切である場合がある 対応できます

1 つ目の欠点に対処するため、最後にもう一つ、 非同期オペコード: kAsyncStartkAsyncUpdatekAsyncDoneアイデア あらゆる 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-startasync-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」オペコード。