XLA（Accelerated Linear Algebra）は、線形代数を最適化する機械学習（ML）コンパイラであり、実行速度とメモリ使用量を改善します。このページでは、XLA コンパイラの目的とアーキテクチャについて簡単に説明します。
目標
現在、XLA はいくつかの ML フレームワーク フロントエンド（PyTorch、TensorFlow、JAX など）をサポートしており、ML 用のオープンソース コンパイラ テクノロジーのエコシステムである OpenXLA プロジェクトの一部となっています。これは、業界をリードする 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 をサポートしています。