このページは Cloud Translation API によって翻訳されました。
Switch to English

XLA用MLIR CodeGen

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

ステージングステップとして、現在、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の移行」のセクションを参照してください。

このプロジェクトでは、LHLOエミッターを可能な限り有効にしたエンドツーエンドの実行可能モデルを優先します。これは、優先度による目標の次の順序リストを意味します。

  • 既存のサンクとエミッターを変更せずに、XLA / GPUをLHLOエミッターで実行可能にします。
  • ケースバイケースで、LHLOのHloInstruction *への参照を削除します。
    • レガシーエミッターをMLIRベースのエミッター(Linalgなど)に切り替える、または
    • 既存のエミッターを機械的に変換してMLIR表現を取得します(GPU方言を使用して標準に移行します)。

サンクの移行(タスク2)

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

  • ホストから呼び出すことができます(xla :: gpu :: Thunk :: ExecuteOnStream())。
  • サブクラスでさまざまなデータを伝送します。
  • BufferAllocation :: SliceおよびStreamExecutorと対話します。
  • カーネルを起動します
  • すべてのランタイムライブラリを呼び出します。

そのコストには以下が含まれます:

  • op固有の構成データ(たたみ込み構成など)を表します。
  • 演算形状とオペランド形状の移行。
  • サンクのツリー(while、conditionなど)を表す。

移行作業は、LHLO /エミッターの移行から独立しています。限られたリソースの下では、LHLO /エミッターの移行の背後で優先されます。

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

  • TFRT
    • (プロ)使用するための素晴らしいCUDAおよびHIPラッパー。
    • (Pro)TFRT opsはC ++コードによって解釈されるため、ライブラリー呼び出し(cuDNN、cuBLAS、cuFFTなど)を簡単に実装できます。
    • (短所)ホスト側は開発中で、テストされていません。
  • Jitted CPUコード
    • (プロ)優れた低能力。いくつかのループと条件を作成すれば完了です。
    • (短所)GPUDialectはまだチェーン/ストリーム/非同期/デバイス割り当てをモデル化していません。
    • (Con)CUDA / HIPランタイムのサポートは最小限です(ツールキットのパス、バージョン、動的ロードなど)。
  • 既存の(解釈)XLAランタイム

決定:TFRTを採用しますが、TFRTでのCPUコードの飛び出しもサポートします。

デバイスLLVM IRの移行(タスク3)

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

ElementalIrEmitterは次の点でユニークです。

  • コードの大部分は、XLA / GPUとCPUの間で共有されます。
  • これは、すべての要素ごとの操作を含む、モデルで見られる操作の大部分を表します。
  • ほとんどのフュージョンはElementalIrEmitterのみに依存しています。
  • op要素とoperand要素間のデータ依存DAGを記述しているため、構造はシンプルです。
  • それは主にポータブルで高レベルです(たとえば、GPU kReduceやGPU kCopyとは異なります)。
  • 動的な形状のサポートは、少なくとも要素単位の操作では簡単です。

さて、すべての操作について、要素的に放出されるかどうかにかかわらず、各XLA操作の最終状態にはいくつかの種類があります。

  1. デバイスコードはLLVM IRのままです。
  2. LHLO-> MLIR LLVM方言のように古いエミッターをリファクタリングします。
    • (コスト)最終的にスタンダードに移行したい場合、使い捨ての作業になります。
    • (メリット)簡単で機械的です。短時間で行えます。
    • (メリット)(1)に比べてメリットがありません。
  3. LHLOのように古いエミッターをリファクタリング-> MLIR GPU +標準+ループ:
    • (コスト)既存のエミッタを標準に引き上げると、いくつかの課題が生じます。ポインターとGEPは、MemRefとサブビューに変換する必要があります。 amdgpuの完全性を保証することも別の問題です。
    • (コスト)XLA / GPUはLLVMメタデータに大きく依存しています。
      • ブロック/スレッドインデックスのrange
      • ロード/ストアのaligndereferenceableinvariant.loadalias.scopenoalias
      • llvm.loop.unroll.disablellvm.loop.unroll.fullllvm.loop.vectorize.enableは順次ループ用です。
    • (メリット)長期的な可能性があります。よりポータブル。
  4. 古いエミッタをLHLO-> Linalgにリファクタリングし、新しいLinalgエミッタを記述します
    • (費用)これはケースバイケースです。以前のオプションと比較して、XLAのパフォーマンスに一致する新しい実装は、ベンチマーク<->最適化ワークフローを通過する必要があります。
    • (メリット)統合スタック。コミュニティのサポート。移植性;より多くの最適化の可能性。

結論:

  • (2)に行かないでください。 (1)または(3)は(2)よりも優れています。 (2)多くの機械的なリファクタリングが必要なため、(1)よりもコストがかかります。 (1)を使用しても、XLAがMLIRエミッターを取得できるようにするという目標を達成できます。これは、LHLO-> LLVM IR->レガシーデバイスエミッターを実行することによって行われます。
  • ElementalIrEmitter opsは(4)に適用されますが、増分は行われません。要素ごとに放出されるすべての操作は同じグラフに接続されるため、操作ごとに操作する方法はありません。この作業は、いくつかの進行中の力(xla / service / mlir_gpu、カーネルジェネレーター、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演算は、そのオペランドバッファ(演算とバッファの間の2つの部分)にのみアクセスできます。 LHLO演算は、オペランド演算にアクセスするために、use-defチェーンを経由する必要があります。
  • ネストされていないレガシーエミッターは、経験的にほとんどオペランドにアクセスしません。唯一の例外はkReduceです。
  • ネストされていないレガシーエミッターは、スライスを取得するためにのみBufferAssignmentにアクセスし、dataflow_analysis()やalias_analysis()などの補助データ構造にはアクセスしません。 llvm_irは、スライス情報に基づいて独自のalias_analysis()を構築します。

結論は、LHLOは大きな手間をかけずに適切に適合すべきであるということです。

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

この手順は、XLAサンクロジックの一部を破棄し始めた場合にのみ必要です(次の手順を参照)。

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

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

MLIRでプロファイリングをサポートするには、おおよそ3つの方法があります。

  • プロファイラーをエンドツーエンドで実行する
  • 挿入されたプロファイラーを使用して、LHLOの各opにプロファイルopを追加します。

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

「注入されたプロファイラ」アプローチには、以下が必要です。

  • LHLOは、プロファイラーをパラメーターとして使用します。
  • 各操作の前後にprofile.start / profile.endを挿入します。
  • これにより、profile。{start、end}がC ++実装に渡されます。

次の理由により、正確なプロファイリングをMLIRで生成されたopsに対して簡単に行うことはできません。

  • MLIRにはタイマーがなく、TFRT / StreamExecutorにも依存しません。
  • MLIRは、複雑なパラメーターを持つC関数を簡単に呼び出すことができません。

ステップ3:(タスク2)サンクの移行

注意点として、サンクにはおおよそ3種類あります。

  • カーネルを起動するKernelThunk。
  • 制御フローサンク。これには、ホスト制御フローロジック(条件付き、while、for、シーケンス)および起動本体カーネルがあります。
  • ライブラリサンク:cuDNN、cuBLAS、cuFFT、NCCLなど

計画は次のとおりです。

  • サンクを(デ)シリアライズ可能にします。
  • TFRTがこれらのセマンティクスをサポートできる状態に改善するのを助けます。
  • 状態が改善したら、個々のサンクを段階的に移行します。

これらのアクションアイテムは、部分的にしか注文されていません。実際の実行順序/エンジニアリングの並列処理は、評価されることになっています。

ステップ4:(タスク3)移行されたElementalIrEmitter

プロファイリングの準備ができたら、MLIR内のすべてのElementalIrEmitterベースのエミッターを完成および調整できます。次に、これらすべてのMLIRベースのエミッターが単一のストリームを使用すると想定して、デフォルトでオンにします。

XLA / CPUのElementalIrEmitterもコードの大部分を共有するため、それらを移行することは有益であることに注意してください。

すべてのベンチマークとパフォーマンスハンティングが完了したら(TODO:パフォーマンスパリティの定義)、新しいMLIRベースのエレメンタルエミッターをオンにし、レガシーElementalIrEmitterを削除します。

この手順では、後で移行するための簡単なフュージョントランジション(ネストされたops)も提供されます。

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

一部のエミッタは、MLIRでサポートするか、機能を削除するまで削除できません。これはMLIRでの比較的大量の作業であり、XLAでのわずかな利益です。マルチストリームXLA / GPUユーザーの現在のユーザーを調査し、妥当な場合はこの機能を削除する必要があります。

ステップ6:(タスク3)移行されたデバイスオペレーション

この手順では、ネストされていないすべてのopsを移行し、ネストされていないすべてのエミッターを削除できます。

これは、kCopyとkReduceのrewrite / refactorを必要とします。 kReduceはすでに多くの作業が行われているため、実行する必要がある実際の作業量はまだ不明です。