XLA 用の MLIR CodeGen

XLA は、HloInstruction で動作し、この表現に多数の最適化を実行して、ターゲットサービス間その多数の最適化を共有します。ある時点で、線形スケジュールが計算され、メモリバッファが統計的に各値に割り当てられます。デバイス固有の codegen は、このシーケンスをトラバースして「エミッター」呼び出すことで、デバイスに最適な表現を生成します（たとえば、CPU での XLA 計算当たりの単一の LLVM 関数、または GPU 演算と GPU をターゲットする際に生成された PTX をカプセル化する一連の「サンク」など）。

ステージング手順として、現在、XLA がバッファー割り当てフェーズを完了し、代わりに lhlo 方言で MLIR モジュールを発行した直後に、プロセスをインターセプトしているところです。ここから、デバイスに応じて MLIR コンポーネント（主に、Linalg、アフィン、GPU 方言）を使用して、codegen を実施します。

次に、codegen 入力として lhlo を使用して XLA/GPU を段階的に移行するための記録の計画を示します。

タスク

ホスト デバイス
入力形式 HloInstruction*（タスク 1） HloInstruction*（タスク 1）
出力形式 xla::Thunk（タスク 2） LLVM IR（タスク 3）
  • タスク 1 は、ホストとデバイス入力形式の両方を HloInstruction* から LHLO に変更します。
  • タスク 2 は、ホストの出力形式をサンクから「ホストのランディングパッド」に変更します（以下を参照）。
  • タスク 3 は、デバイス出力を LLVM IR から何らかの MLIR に移行します。このプロジェクトではオプションです。詳細は、「デバイス LLVM IR を移行する」セクションをご覧ください。

| ホスト | デバイス --- | --- | --- 入力形式 | HloInstruction（タスク 1） | HloInstruction（タスク 1） 出力形式 | xla::Thunk（タスク 2） | LLVM IR（タスク 3）

  • 既存のサンクとエミッターは変更せずに、LHLO エミッタ―で XLA/GPU を実行可能にする。
  • LHLO の HloInstruction* への参照をケースに応じて排除する。
    • レガシーエミッターを MLIR ベースのエミッター（Linalg など）に切り替える。
    • MLIR 表現を取るように既存のエミッターを機械的に変換する（GPU 方言を使った標準に移行する）。

サンクを移行する（タスク 2）

xla::gpu::Thunk は、次のようなデータ構造です。

  • ホストから呼び込める（xla::gpu::Thunk::ExecuteOnStream()）。
  • サブクラスでさまざまなデータを運搬する。
  • BufferAllocation::Slice と StreamExecutor と対話する。
  • カーネルを起動する。
  • すべてのランタイムライブラリに呼び出す。

次のようなコストがあります。

  • 演算固有の構成データ（畳み込み構成など）を表現する。
  • 演算の形状とオペランドの形状を移行する。
  • サンクのツリーを表現する（while 句、条件など）。

移行作業は、LHLO/エミッターの移行とは独立して行われます。限られたリソースの中で、LHLO/エミッター移行の後に優先して行われます。

ホスト側のパーツを LHLO から下げる方法にはいくつかの選択肢があります。

  • TFRT
    • （メリット）優れた CUDA および HIP ラッパーの使用。
    • （メリット）TFRT 演算は C++ コードで解釈されるため、ライブラリ呼び出し（cuDNN、cuBLAS、cuFFT など）を実装しやすい。
    • （デメリット）ホスト側は開発中であり、検証されていない。
  • JIT コンパイルされた CPU コード
    • （メリット）優れた低位機能。ループと条件をいくつか作成すれば、完了です。
    • （デメリット）GPUDialect は、まだチェーン/ストリーム/非同期/デバイス割り当てをモデル化しない。
    • （デメリット）最小限の CUDA/HIP ランタイムサポート（ツールキットパス、バージョン、動的読み込みなど）。
  • 既存の（解釈）XLA ランタイム

判定: TFRTを採用しますが、TFRT での Jit コンパイルする CPUコードもサポートします。

デバイス LLVM IR を移行する（タスク 3）

エレメンタルエミッターは、要素ごとに埋めることでターゲット演算を生成します。 各出力要素は、オペランドの要素のセットに依存します。すべての要素は、バッファを動的インデックスと組み合わせることによって記述されます。ほとんどすべての「数学」演算を記述するだけで十分ですが、パフォーマンス上の理由から、「数学」演算の大きなサブセットのみが（Cpu|Gpu）ElementalIrEmitter に直接実装されます。

ElementalIrEmitter は、次の点で特異です。

  • コードの大部分が XLA/GPU と CPU で共有されている。
  • すべての要素ごとの演算を含め、モデルに見られる大部分の演算を表現する。
  • ほとんどの融合は ElementalIrEmitter のみに依存する。
  • 演算要素とオペランド要素間のデータ依存関係 DAG を記述するため、単純な構造を持つ。
  • ほぼポータブルで高レベルである（GPU kReduce と GPU kCopy など）。
  • 少なくとも要素ごとの演算で、動的な形状のサポートを簡単に行える。

要素的にエミットされているかどうかにかかわらず、すべての演算において、各 XLA 演算子の終了状態にはいくつかのフレーバーがあります。

  1. デバイスコードを LLVM IR のままとする。
  2. 古いエミッターを LHLO -> MLIR LLVM 方言となるようにリファクタリングする。
    • （コスト）最終的に標準に移行する場合に使い捨て作業となる。
    • （メリット）簡単で機械的である。短期間で実行可能。
    • （メリット）（1）に比べるとそれ以上のメリットはない。
  3. 古いエミッターを LHLO -> MLIR GPU + 標準 + ループとなるようにリファクタリングする。
    • （コスト）既存のエミッターを標準にリフトすると、問題が発生する。ポインタと GEP は、MemRefs と SubViews に変換する必要がある。amdgpu の完全性を保証することももう 1 つの問題。
    • （コスト）XLA/GPU は LLVM メタデータに大きく依存している。
      • range: インデックスのブロックまたはスレッド化
      • align、dereferenceable、invariant.load、alias.scope、noalias: 読み込み/格納
      • llvm.loop.unroll.disable、llvm.loop.unroll.full、llvm.loop.vectorize.enable: シーケンシャルループ
    • （メリット）長期的に使用可能。よりポータブル。
  4. 古いエミッターを LHLO -> Linalg となるようにリファクタリングし、新しい Linalg エミッターを記述する。
    • （コスト）これはケースによる。前のオプションと比べると、XLA のパフォーマンスに一致する新しい実装は、ベンチマーク <-> 最適化のワークフローを通過する必要があり、演算によっては著しいコストとなる可能性がある。
    • （メリット）統合スタック、コミュニティサポート、移植性、より多くの最適化の可能性

結論:

  • （2）は使用しません。（1）または（3）が（2）より優れています。（2）は機械的なリファクタリングを多く要するため、（1）よりもコストがかかります。（1）を採用した場合、XLA で MLIR エミッターを拾う目標を達成することができます。これは、LHLO -> LLVM IR -> レガシーデバイスのエミッターの実行によって行います。
  • ElementalIrEmitter 演算は（4）に適用されますが、増分には適用されません。要素でエミットされる演算は同じグラフに接続されるため、演算ごとに実行する術がないからです。この作業は複数の継続する力（カーネルジェネレータ、Linalg）の統一ポイントとしても機能します。
  • ほかのすべての演算は（1）に適用されます。拡張目標として、これらは（3）または（4）に移行できる可能性があります。

優先

前述の 3 つのタスクは並列化可能ですが、限定されるリソースの中でシリアル化される必要があります。優先順位は、各タスクの目に見える完了結果に焦点が当てられます。

優先順位は、タスク 1（レガシーエミッターの LHLO）> タスク 2（サンク）> タスク 3（MLIR エミッター）です。

タスク 1 の終了までに、XLA のユーザーは LHLO（カーネルジェネレータ）を生成して実行できます。コンパイル形式はシリアル化可能な MLIR ではありません。

タスク 2 の終了までに、LHLO は適切なシリアル化可能な MLIR に低下します。これにより、オフラインコンパイルが可能となります。

タスク 3 の終了までに、すべての XLA エミッターは実装において MLIR ベースとなります。

詳細な設計

ステップ 1:（タスク 1）LHLO を完了してレガシーエミッターに LHLO を取らせる

このステップでは、既存のすべての XLA/GPU エミッターを MLIR 演算と対話させます。このステップは純粋なリファクタリングと NFC です。

このステップはほぼ機械的に行われますが、ネストされていない HloComputation と LHLO には次のような違いがあることに注意してください。

  • 各 HloInstruction は、オペランド（データフロー DAG）に直接アクセスできます。一方、各 LHLO 演算は、オペランドのバッファ（演算とバッファ間の二部）にしかアクセスできません。LHLO 演算は use-def チェーンを経由してオペランド演算にアクセスする必要があります。
  • ネストされていないレガシーエミッターは、経験的に決してオペランドにアクセスしません。唯一の例外は kReduce です。
  • ネストされていないレガシーエミッターは、dataflow_analysis() や alias_analysis() のような補助データ構造にアクセスするためではなく、スライスを取得するためにのみ BufferAssignment にアクセスします。llvm_ir は、スライス情報に基づいて独自の alias_analysis() を構築します。

結論としては、LHLO が大きな手間をかけることなく適合すると言えます。

ステップ 2:（オプション）プロファイリングのサポート

このステップは、いくつかの XLA サンク論理（次のステップを参照）を破棄し始める場合にのみ必要です。

MLIR ベースのエミッターを実際に有効にする前に、MLIR ベースエミッターのプロファイリングが必要です。

現在 XLA は、StreamExecutor のタイマーに呼び出すことで独自のプロファイリングを実行します。このタイマーは内部的に、カーネルの起動前後に 2 つのイベントを挿入し、その 2 つのイベント間の同期時間を測定します。

MLIR でプロファイリングをサポートするには、大きく 3 つのアプローチがあります。

  • エンドツーエンドでプロファイラを実行する
  • LHLO の各演算に対し、注入されたプロファイラを使ってプロファイル演算を追加する

「エンドツーエンド」のアプローチは MLIR に透過的ですが、XLA がそもそもそれを使用しないのと同じ問題があります。プロファイラから収集されるライブラリの呼び出し（nvprof/...）は HLO 演算に簡単に関連付けられないのです。たとえば、cuDNN は各 HLO に複数のカーネルを起動するため、どのカーネルがどの HLO に対応するのかを見分けることが困難です。

「注入されたプロファイラ」アプローチには、次の項目が必要となります。

  • LHLO がプロファイラをパラメータとして取ること
  • 各演算の前後に profile.start / profile.end を挿入する必要があること
  • profile.{start,end} から C++ 実装に低下させるパス

実際のプロファイリングは、次の理由により MLIR が生成した演算では簡単に実行できません。

  • MLIR にはタイマーがなく、TFRT / StreamExecutor に依存しないため
  • MLIR は、複雑なパラメータを使用する C 関数に簡単に呼び出さないため

ステップ 3:（タスク 2）サンクを移行する

注意書きとして、サンクには大きく 3 つの種類があります。

  • KernelThunk: カーネルを起動します。
  • 制御フローサンク: ホスト制御フロー論理（条件、While、For、シーケンス）があり、ボディカーネルを起動します。
  • ライブラリサンク: cuDNN、cuBLAS、cuFFT、NCCL など

計画は次のとおりです。

  • サンクを（逆）シリアル化可能にする
  • TFRT をこれらのセマンティックをサポートできる状態に改善する
  • 状態が改善されたら、個別のサンクを徐々に移行する

上記のアクションアイテムは一部のみ順序付けられています。実際の実行順/エンジニアリング並列処理は、進行中に評価されます。

ステップ 4:（タスク 3）移行済みの ElementalIrEmitter

プロファイリングの準備ができたら、MLIR のすべての ElementalIrEmitter ベースエミッターを完了して調整することができます。その後で、すべての MLIR ベースのエミッターが単一のストリームを使用することを前提に、ElementalIrEmitter ベースエミッターをデフォルトで有効にします。

コードの大部分を共有する XLA/CPU の ElementalIrEmitter を移行することにもメリットがあります。

すべてのベンチマーク設置とパフォーマンスの調整が完了したら（TODO: パフォーマンスパリティの定義）、新しい MLIR ベースのエレメンタルエミッターを有効にし、レガシー ElementalIrEmitter を削除します。

このステップでは、後の移行で融合への遷移（ネスト演算）を簡単に行えるようにします。

ステップ 5: マルチストリームのサポートまたは排除

MLIR でマルチストリームがサポートされるまで、またはその機能を排除するまで、一部のエミッターを削除できません。MLIR では比較的大量の作業であっても、XLA のゲインはわずかです。マルチストリーム XLA/GPU ユーザーの現在のユーザー数を調べ、合理的であればこの機能を削除するようにする必要があります。

ステップ 6:（タスク 3）移行済みのデバイス演算

このステップはネストされていないすべての演算を移行し、その後でネストされていないエミッターを削除することができます。

これには、kCopy と kReduce のリライト/リファクタリングが必要です。kReduce はすでに多数の作業を完了済みであるため、どれほどの作業が残っているのかを確認する必要があります。