XLA (Accelerated Linear Algebra)（エックスエルエー）
Google が開発した線形代数特化のドメイン固有コンパイラ。TensorFlow と JAX のデフォルトコンパイラとして、計算グラフ全体を最適化し TPU・GPU・CPU 向けの効率的なマシンコードを生成する。演算融合とメモリ最適化により大規模モデルの学習・推論を高速化する。
XLA とは
XLA（Accelerated Linear Algebra）は Google が 2017 年に公開したドメイン固有コンパイラで、TensorFlow と JAX のバックエンドとして計算グラフをネイティブコードに変換する。線形代数演算（行列積・畳み込み・要素単位演算）を対象とし、計算グラフ全体を俯瞰した最適化（演算融合・バッファ再利用・レイアウト最適化）を実行する。
コンパイルパイプライン
XLA のコンパイルは以下の段階を経る。
| ステージ | 処理内容 |
|---|---|
| HLO 生成 | フレームワーク（TF/JAX）から High Level Optimizer IR を受け取る |
| HLO 最適化 | 演算融合・定数畳み込み・代数簡約・CSE（共通部分式除去） |
| バッファ割当 | テンソルのメモリレイアウトとライフタイム解析・バッファ再利用 |
| コード生成 | ターゲット固有のコード生成（TPU: TPU ISA / GPU: PTX via LLVM / CPU: x86/ARM via LLVM） |
XLA の主要最適化
- 演算融合（Op Fusion）: element-wise 演算の連鎖を単一カーネルに融合。メモリ帯域の消費を 80% 以上削減するケースもある
- メモリレイアウト最適化: テンソルの物理配置をハードウェアのメモリ階層に最適化（例: TPU の 128x128 タイル配置）
- 自動シャーディング: GSPMD（General and Scalable Parallelization for ML Computation Graphs）により、複数デバイスへの演算分割を自動化
- リマテリアライゼーション: メモリ不足時に中間テンソルを再計算で補い、ピークメモリを削減
TPU との関係
XLA は元々 Google の TPU（Tensor Processing Unit）向けコンパイラとして開発された。TPU のハードウェア設計と XLA のコンパイラ最適化は密結合しており、以下のような TPU 固有の最適化が含まれる。
- MXU マッピング: 128x128 の行列乗算ユニットに演算を効率的にマッピング
- HBM 管理: TPU v4 の 32GB HBM をコンパイラが自動管理
- ICI 通信最適化: TPU Pod 内のチップ間インターコネクト通信のスケジューリング
JAX との統合
JAX は XLA をデフォルトコンパイラとして使用し、jax.jit デコレータで Python 関数を XLA コンパイル対象にする。
jax.jit: 関数単位の JIT コンパイルjax.pmap: 複数デバイスへの並列マッピング（内部で XLA の SPMD を使用）jax.vmap: バッチ次元の自動ベクトル化
JAX + XLA の組合せは Google の大規模言語モデル学習（PaLM、Gemini）の基盤技術として使用されている。
FAQ
Q1: XLA は PyTorch でも使えるか？
PyTorch/XLA プロジェクトにより、PyTorch モデルを XLA 経由で TPU や GPU 上で実行可能。ただし PyTorch 2.x では torch.compile + TorchInductor が主流であり、PyTorch/XLA は主に TPU 利用時に選択される。
Q2: XLA と TorchInductor の違いは？
XLA はグラフ全体を一括コンパイルする AOT（Ahead-of-Time）方式、TorchInductor は部分グラフ単位で JIT コンパイルする方式。XLA は TPU で最高性能を発揮するが、GPU では TorchInductor + Triton が同等以上の性能を出すケースが増えている。
Q3: XLA のコンパイル時間は問題にならないか？
初回コンパイルに数分〜数十分かかる場合がある（大規模モデルほど長い）。JAX では jax.jit のキャッシュにより 2 回目以降は高速化される。また、AOT コンパイル結果をディスクに保存して再利用する機能もサポートされている。