HLO から Thunk へ

このドキュメントでは、XLA High Level Optimizer（HLO）モジュールの初期状態から最終的な実行可能ファイルまでの流れについて説明します。「モジュール」を省略して「HLO」と呼ぶこともあります。

HLO から thunk への図

最適化前の HLO

まず、最適化前の HLO モジュールから始めます。最適化前の HLO には、fusion や bitcast など、XLA の内部と見なされるオペレーション（ops）は含まれません。この段階では、Ops にレイアウトがないか、レイアウトがあっても無視されます。最適化前の HLO は通常、TensorFlow や JAX などの高レベルのフレームワークによって生成されます。XLA フラグ -xla_dump_to を使用すると、最適化前の HLO がファイル名接尾辞「before_optimizations.txt」のファイルにダンプされます。

HLO モジュールを最適化する

XLA:GPU パイプラインは、一連のパスを実行して、最適化前の HLO を最適化された HLO に変換します。パスは意味的にグループ化され、次の順序で実行されます。

これには、Shardy Partitioner などのパスや、SPMD シャーディング用のパスが含まれます。

最適化パス

これには、正当化パスと簡素化パスの両方が含まれます。

集団最適化パス

最適化パスと似ていますが、集合演算に重点を置いています。

レイアウト割り当てパス

各 HLO オペレーションには、命令形状の一部であるレイアウトが割り当てられます。レイアウトは、テンソルがメモリに物理的に配置される方法を制御します。

レイアウトを含むシェイプの例:

f32[10,20,30]{2,0,1}

要素型の後に、シェイプの論理ディメンションが続き、レイアウト順列がマイナーからメジャーの順に続きます。この例では、最もマイナーなディメンションは 30、2 番目にマイナーなディメンションは 10、メジャー ディメンションは 20 です。

レイアウト割り当ての目的は、貪欲法を使用して必要な物理転置の数を最小限に抑えることです。特定のレイアウト制約（cuDNN/cuBLAS ライブラリは連続したディメンションを想定するなど）から始まり、レイアウトを HLO グラフの「下」に伝播してから「上」に伝播します。レイアウト伝播の終了時に、オペランドから伝播されたレイアウトとユーザーから伝播されたレイアウトが競合する場合があります。この競合を解決するために、オペランド レイアウトから命令レイアウトにレイアウトを変更する copy HLO 命令が挿入されます。

レイアウト正規化パス

物理的な形状を把握することはやや難しいため、レイアウトの正規化では、デフォルトのレイアウト {rank-1, rank-2, …, 0} を使用するように形状を書き換えようとします。上記の例では、正規化された形状は f32[20,10,30]{2,1,0} になります。レイアウトを変更するコピー オペレーションは、transpose と bitcast の組み合わせとして書き換えられます。現在、すべてのオペレーションを正規化することはできないため、デフォルト以外のレイアウトを持つオペレーションがまだいくつかあります。特に gather と dot がそうです。正規化されたオペレーションと正規化されていないオペレーションの境界には、転置を表す bitcast オペレーションがあります。つまり、レイアウトが割り当てられた転置は、物理的には no-op になります。

レイアウトの正規化では、暗黙的な転置を明示的にすることもできます。これは、codegen が専用のエミッタで明示的な転置を処理できるため重要です。たとえば、オペランドと結果の間で異なる物理レイアウトを持つことは、技術的には許可されています（ランクが異なる場合など）。レイアウト正規化パスの一部として実行される ReshapeDecomposer パスは、リシェイプを transpose、リシェイプ bitcast、transpose のシーケンスに変換します。

レイアウト割り当ての最適化パスを投稿する

ここで最も重要なパスは、Triton フュージョン（GEMM フュージョン + Softmax/Layernorm フュージョン）またはライブラリ呼び出しへの書き換えです。このステップでは自動チューニングも実行されます。XLA はさまざまなエミッタから選択し、畳み込みまたはドットに最適なアルゴリズムを選択し、Triton エミッタによって処理されるフュージョンに最適なタイリングを見つけます。

Fusion パス

主なパスは PriorityFusion と Multi-Output の融合です。

PriorityFusion では、コストモデルに基づいてフュージョンを形成します。融合時に、オペレーションをすべてのユーザーに融合できる場合は、複数のユーザーでオペレーションを複製できます。また、可能であれば、既存の Triton Softmax フュージョンを拡張することも許可します。

Multi-Output フュージョンは、オペランドを共有するオペレーション/フュージョンを融合できる別のパスです。また、オペランド/オペランド融合を重複なしでユーザーに融合することもできます。その場合は、追加の出力が追加されるため、融合される op の他のユーザーをこれらの出力にリダイレクトできます。このパスでは、HLO グラフにサイクルが導入されないように注意する必要があります。

Multi-Output 融合の後、共通部分式の削除（HloCSE パス）が実行され、以前に重複したオペレーションが同じ融合に収まった場合、それらが再び統合される可能性があります。

複数の融合後パス

コレクティブに関連する複数のパス（コレクティブを非同期に変換したり、コレクティブの特定の相対順序を適用したりするなど）。

最後に CopyInsertion を実行してコピーを追加し、インプレース オペレーションで他の場所でまだ必要なデータが上書きされないようにします。

最適化の最後に、フラグ -xla_dump_to を使用している場合、最適化された HLO はファイル名接尾辞「after_optimizations.txt」を持つファイルにダンプされます。HloModule を実際に変更する中間パスの後に HLO をダンプする場合は、-xla_dump_hlo_pass_re=.* フラグ（または特定の正規表現を使用して特定のパスに制限）を使用できます。

スケジュール

スケジュールがない HLO モジュールでも、オペレーションが処理される順序にはある程度の自由度があります。オペランドと結果の関係および制御依存関係を尊重するトポロジカル ソートはすべて有効です。スケジューリングによって、使用する特定の順序が決まります。この段階での主な懸念事項は、テンソルのライフタイムに依存する最大メモリ消費量です。最初のステップでは、さまざまなスケジューラ アルゴリズムを試して、ピーク時のメモリ消費量を最小限に抑えるスケジュールを選択します。この時点では、まだ物理バッファを扱っておらず（「バッファの割り当て」で扱います）、メモリ使用量をシミュレートしていることに注意してください。

次に、LatencyHidingScheduler パスが実行され、コンピューティングと通信のオーバーラップを最大化しようとします。ただし、メモリ使用量が再び増加する可能性があります。

最後に、メモリ使用量のピークが使用可能なメモリ量を超えている場合は、HloRematerialization を実行します。このパスでは、パフォーマンスを犠牲にしてメモリ使用量を削減しようとします。たとえば、一部のフュージョンが分割され、一部のオペレーションが複製されてバッファのライフタイムが短縮されることがあります。再実体化が発生した場合は、モデル側のメモリ要件を減らす方法（バッチサイズを小さくするなど）を検討することをおすすめします。

バッファの割り当て

LLVM IR に変換する直前に、バッファ スライスを HLO グラフの各命令に割り当てるバッファ割り当てパスを実行します。バッファの割り当ては、次の手順で行われます。

  1. HloDataflowAnalysis は、HloValues（基本的に論理バッファ）を命令に割り当てます。インプレース オペレーションの場合、オペランドの HloValue を再利用できます。op は複数の HloValue を定義できます（タプル結果の形状など）。

  2. HloAliasAnalysis は、エイリアシング オペレーションのバッファを結合しようとし、HloValue から HloBuffer へのマッピングを計算します。

  3. BufferAssignment は、同じバッファ スライスがライフタイムが重複する異なる HloBuffers に使用されないように、HloBuffers のマッピングを大きなバッファ内のバッファ スライスに計算します。エイリアス化される可能性があるオペレーションについては、わずかな重複があっても問題ありません（一方の HloBuffer の終了時刻が他方の HloBuffer の開始時刻と一致する可能性があります）。フラグ -xla_dump_to を使用すると、バッファ割り当てに関する一部の情報が「after_optimizations-buffer-assignment.txt」という名前の接尾辞が付いたファイルにダンプされます。

Thunk

HLO グラフが最適化され、スケジュールされると、特定のバックエンド（CPU または GPU）のサンクの線形シーケンスに変換されます。

XLA では、Thunk は、ランタイムが実行する自己完結型の作業単位の抽象化です。コンパイル済みカーネルの起動、特定のオペレーション、ライブラリ呼び出し、制御フロー構造、集団通信などがあります。Thunk Sequence は、特定のバックエンドの実行可能ファイル全体を表します。

Thunk Emission

スケジュールされた HLO 計算をサンク シーケンスに変換するプロセスを「サンク エミッション」と呼びます。これは、各バックエンドの専用のエミッタクラスによって処理されます。

GPU バックエンドの場合、これは IrEmitterUnnested によって処理されます。EmitHloComputation は、計算内の HLO 命令のスケジュールされたリストを反復処理し、専用の Emit... メソッド（EmitFusion、EmitConvolutionThunk、EmitWhile）。これらの各メソッドは、適切な Thunk オブジェクトを構築し、Thunk シーケンスに追加します。

CPU バックエンドの場合、ThunkEmitter がこの役割を果たし、同様の方法で整理されます。最終的な ThunkSequence は CpuExecutable に埋め込まれます。

HLO モジュールのエントリ計算の各命令は、最終的なサンク シーケンスのサンク（kTuple、kConstant など）に対応しない場合もあれば、1 つまたは複数（並べ替え命令など）に対応する場合もあります。

コマンド バッファ: GPU での実行の最適化

最新の GPU ハードウェアでは、GPU オペレーション（カーネル起動、メモリコピーなど）のシーケンスを一度記録し、CPU オーバーヘッドを最小限に抑えてシーケンスを複数回再生できます。これは、特に多くの小さな高速起動カーネルを含むワークロードにとって重要なパフォーマンス最適化です。XLA は、CUDA グラフまたは HIP グラフの抽象化としてコマンド バッファを使用します。コア インターフェースは GpuCommandBuffer で定義されます。

コマンド バッファは、サンク シーケンスで CommandBufferThunk によって表されます。

エミッタは、このサンクを HLO 命令から直接生成しません。代わりに、ThunkSequence 自体で実行される CommandBufferConversionPass によって行われます。

このパスは、互換性のあるサンクの連続するサブシーケンス（KernelThunk と GemmThunk のシーケンスなど）を識別します。次に、見つかったサブシーケンスを単一の CommandBufferThunk に置き換えます。新しいサンクは、元のサンクのロジックを軽量の CommandBufferCmd オブジェクトのリストとしてカプセル化します。特定の GPU ストリームで CommandBufferThunk が初めて実行されると、コマンドのシーケンスがハードウェア コマンド バッファに「記録」されます。以降の実行では、記録されたシーケンスを「再生」する単一のコマンドを GPU に発行するだけです。これにより、個々のカーネルを起動する CPU オーバーヘッドを回避できます。

実行可能

XLA コンパイル パイプラインの最終的な成果物は、自己完結型のプラットフォーム固有の実行可能ファイルです。このオブジェクトは、CPU や GPU などのターゲット デバイスでコンパイル済みプログラムを実行するために必要なすべての情報をカプセル化します。コンパイラとランタイムの橋渡しをするものです。PJRT などの最新のランタイムは、やや上位レベルの抽象化を使用しますが（PjRtExecutable を参照）、最終的にはバックエンド固有の実行可能ファイルをラップします。

Executable には、コンパイル中に生成された重要な情報がいくつか含まれています。具体的な内容はバックエンドによって異なりますが、一般的には次のものが含まれます。

  • コンパイル済みコード: デバイスで実行される低レベルのマシンコードです。CPU の場合、通常は 1 つ以上のオブジェクト ファイルです。GPU の場合、これは PTX または HSACO 形式でコンパイルされたデバイスコードであり、実行時に GPU に読み込まれます。

  • 実行プラン（ThunkSequence）: ランタイム ロジックの中核。これは Thunk オブジェクトの線形シーケンスです。各サンクは、カーネルの起動、ライブラリ関数（cuBLAS など）の呼び出し、制御フローの処理などの単一の作業単位を表します。ランタイムは、このシーケンスを反復処理してプログラムを実行します。

  • メモリ レイアウト（BufferAssignment）: BufferAssigner によって生成されるこの重要なメタデータは、計算の完全なメモリ レイアウトを表します。各バッファのサイズと、パラメータ、出力、一時値に対してメモリが割り当てられ、再利用される方法を指定します。ランタイムはこれを使用して、デバイスメモリを割り当て、各サンクに正しいポインタを渡します。

  • （省略可）HLO モジュール: デバッグとプロファイリングのために、実行可能ファイルは、コンパイル元の最終的な最適化された HloModule への参照を保持することがよくあります。

最終的な実行可能ファイルの作成は、特定のバックエンドごとにコンパイラによってオーケストレートされます。コンパイラ実装の RunBackend メソッドは、コンパイル プロセスの最終ステップであり、コンパイルされたすべてのアーティファクトを Executable オブジェクトにパッケージ化します。GpuCompiler と CpuCompiler は、それぞれ GPU と CPU をターゲットにします。

ユーザーが実行可能ファイルで Execute... を呼び出すと、ランタイムは BufferAssignment を使用してメモリを割り当て、ThunkSequence を呼び出してコンパイル済みコードを使用してデバイスでオペレーションを開始します。