XLA でのエイリアス設定

このドキュメントでは、XLA エイリアス API について説明します。この API を使用すると、XLA プログラムの作成時に入力バッファと出力バッファ間のエイリアスを指定できます。

コンパイル時のエイリアスの定義

たとえば、入力に 1 を追加するだけの簡単な HLO モジュールについて考えてみましょう。

HloModule increment

ENTRY entry {
  %p = f32[] parameter(0)
  %c = f32[] constant(1)
  ROOT %out = f32[] add(%p, %c)
}

このモジュールは 2 つの 4 バイトバッファを割り当てます。1 つは入力 %p 用、もう 1 つは出力 %out 用です。

ただし、多くの場合、インプレースで更新を行うことが望まれます(たとえば、式を生成するフロントエンドで、インクリメント p++ のように、計算後に入力変数が存在しなくなった場合など)。

このような更新を効率的に実行するには、入力エイリアスを指定します。

HloModule increment, input_output_alias={ {}: 0 }

ENTRY entry {
  %p = f32[] parameter(0)
  %c = f32[] constant(1)
  ROOT %out = f32[] add(%p, %c)
}

この形式では、出力全体({} でマーク)が入力パラメータ 0 にエイリアスされることを指定します。

エイリアス設定をプログラムで指定するには、XlaBuilder::SetUpAlias API をご覧ください。

実行時にエイリアスを定義する

前の手順で定義したエイリアスは、コンパイル時に指定されます。 実行中、LocalClient::RunAsync API を使用して、バッファを提供するかどうかを選択できます。

プログラムへの入力バッファは ExecutionInput でラップされ、MaybeOwningDeviceMemory のツリーを含みます。メモリが「所有」として指定されている場合(バッファの所有権が XLA ランタイムに渡される)、バッファが実際に提供され、コンパイル時エイリアス API がリクエストするように、その場で更新が実行されます。

ただし、コンパイル時にエイリアスされたバッファが実行時に寄付されない場合、コピー保護が機能します。つまり、追加の出力バッファ O が割り当てられ、エイリアスされるはずの入力バッファ P の内容が O にコピーされます。これにより、プログラムは実質的に、バッファ O が実行時に提供されたかのように実行できます。