XLA（Accelerated Linear Algebra）は、線形代数を最適化して実行速度とメモリ使用量を改善する機械学習（ML）コンパイラです。このページでは、XLA コンパイラの目的とアーキテクチャの概要について説明します。
目標
現在、XLA は複数の ML フレームワーク フロントエンド（PyTorch、TensorFlow、JAX など）をサポートしており、OpenXLA プロジェクトの一部になっています。OpenXLA プロジェクトは、ML ハードウェアとソフトウェアの大手組織が共同で開発した ML 用のオープンソース コンパイラ テクノロジーのエコシステムです。OpenXLA プロジェクトが作成される前は、TensorFlow プロジェクト内で XLA が開発されましたが、基本的な目的は変わりません。
実行速度を向上させる。サブグラフをコンパイルして有効期間の短い演算の実行時間を短縮し、ランタイムのオーバーヘッドを排除します。また、パイプライン オペレーションを融合してメモリのオーバーヘッドを削減します。また、既知のテンソル形状を特殊化して、より積極的な定数伝播を可能にします。
メモリ使用量を改善する。メモリ使用量を分析してスケジュールし、多くの中間ストレージ バッファを排除します。
カスタム オペレーションへの依存を減らす。手動で融合されたカスタム オペレーションのパフォーマンスに合わせて、自動的に融合された低レベルのオペレーションのパフォーマンスを向上させることで、多くのカスタム オペレーションが不要になります。
ポータビリティを改善する。新しいハードウェア用の新しいバックエンドを比較的簡単に記述し、そのハードウェアで ML モデルの大部分を変更せずに実行できるようにします。これは、新しいハードウェア用に個別のモノリシック演算に特化したアプローチとは対照的です。この場合、これらの演算を利用するにはモデルの書き換えが必要になります。
仕組み
XLA コンパイラは、StableHLO で定義された ML フレームワークからモデルグラフを取得し、さまざまなアーキテクチャのマシン命令にコンパイルします。StableHLO は、ML フレームワークとコンパイラの間にポータビリティ レイヤを提供する、バージョニングされたオペレーション セット（HLO = 高レベルのオペレーション）を定義します。
一般に、モデルグラフをターゲットに最適化された実行可能ファイルに変換するコンパイル プロセスには、次の手順が含まれます。
XLA は、CSE、ターゲットに依存しない演算融合、計算用のランタイム メモリを割り当てるためのバッファ分析など、ターゲットに依存しないいくつかの組み込み最適化および分析パスを、StableHLO グラフに対して実行します。この最適化ステージでは、XLA によって StableHLO 言語も内部 HLO 言語に変換されます。
XLA は HLO 計算をバックエンドに送信して、さらに HLO レベルの最適化を行います。今回はターゲット固有の情報とニーズを念頭に置いて行います。たとえば、GPU バックエンドは、特に GPU プログラミング モデルに有用な演算融合を実行し、計算をストリームに分割する方法を決定できます。この段階で、バックエンドは特定のオペレーションまたはその組み合わせにパターンマッチングを行い、最適化されたライブラリ呼び出しを行う場合もあります。
その後、バックエンドはターゲット固有のコード生成を実行します。XLA に含まれる CPU と GPU のバックエンドは、低レベルの IR、最適化、コード生成に LLVM を使用します。これらのバックエンドは、HLO の計算を効率的に表すために必要な LLVM IR を出力し、次に LLVM を呼び出してこの LLVM IR からネイティブ コードを出力します。
このプロセス内で、XLA コンパイラは、新しい HW アーキテクチャをターゲットとするために代替バックエンドに簡単にスロットリングできるという点で、モジュール式です。GPU バックエンドは現在、LLVM NVPTX バックエンドを介して NVIDIA GPU をサポートしています。CPU バックエンドは複数の CPU ISA をサポートしています。