Join us at TensorFlow World, Oct 28-31. Use code TF20 for 20% off select passes. Register now

XLA 概要

XLA(Accelerated Linear Algebra)は線形代数の演算に特化したコンパイラで、XLAを使うことでTensorFlowの演算を最適化し、メモリ使用量、性能、サーバやモバイル環境での移植性の面での改善が期待できます。現在のところ、ほとんどのユーザにとってXLAを使うことによる大きな恩恵は得られないかもしれませんが、just-in-time (JIT) コンパイル機能ahead-of-time (AOT) コンパイル機能 を通して、実験的にXLAを使っていただくのは歓迎です。ただし、新たなハードウエアアクセラレータの開発者については、XLAを試してみることをおすすめします。

XLAは実験的なフレームワークであり、現在活発に開発されています。既存のオペレーションのセマンティクスが変わることはほとんどないと思いますが、重要なユースケースに対応するために新たなオペレーションが追加されることがあります。XLAチームは、GitHubを通したコミュニティへの貢献や、不足している機能に関するフィードバックを歓迎します。

なぜXLAを開発したか?

TensorFlowと連携するXLAを開発した目的はいくつかあります。

  • 実行速度の向上: サブグラフをコンパイルし、TensorFlowランタイムのオーバヘッドを削減することで、軽量なオペレーションの実行時間を短縮します。また、複数オペレーションを結合することで、メモリのオーバヘッドを削減します。さらに、既知のテンソルの形状に特化することで、より積極的に定数畳み込みができるようにします。
  • メモリ使用量の改善: メモリ使用量の解析とスケジューリングによって、オペレーションの中間データを保持する領域を削減します。
  • 独自オペレーションへの依存度削減: 自動的に結合された低レベルなオペレーションの性能を向上させ、人手による独自オペレーションと同等の性能を得られるようにすることで、独自オペレーションを作成する必要性が無くなります。
  • モバイル環境でのディスク占有スペース削減: サブグラフをAOTコンパイルし、他のアプリケーションに直接リンク可能なオブジェクトとヘッダを出力することで、TensorFlowのランタイムを削除します。これによって、モバイルでの推論時のディスク占有スペースを桁違いに削減できます。
  • 移植性の向上: TensorFlowのプログラムの大部分を修正することなく、新しいハードウェア向けのバックエンドを書くことが容易になります。これは、TensorFlowのプログラムを書き換えて新しいハードウェア向けのオペレーションを作る方式とは全く異なるものです。

XLAはどのように動くのか?

XLAの入力となる言語は、"HLO IR"または単にHLO(High Level Optimizer)と呼ばれます。HLOのセマンティクスは、Operation Semantics に記載されています。HLOは、コンパイラで扱う 中間表現 と考えるとわかりやすいかもしれません。

XLAはHLOで定義されたグラフ("Computations")を、様々なハードウェアアーキテクチャの実行コードにコンパイルします。Developing a new backend for XLA に記載されているように、XLAを新たなハードウェアで動作させることが容易であるという点で、XLAはモジュール化されていると言えます。CPU (x86-64、ARM64) 向けの処理NVIDIA GPU向けの処理 は、TensorFlowのメインのソースコードツリーから参照することができます。

次に示す図は、XLAの内部で行われているコンパイル処理を示しています。

XLAは、CSE やオペレーション結合、計算時のメモリ割り当て解析といった、ターゲットに依存しない最適化や解析を行います。

ターゲットに依存しない最適化処理の後、XLAはHLO computationをバックエンドに転送します。バックエンドでは、ターゲット固有の情報を考慮してHLOレベルでの最適化を行います。例えば、バックエンドがXLA GPUである場合、よりGPUのプログラミングモデルに適したオペレーションの結合に加え、computationを複数のストリームへ分割して割り当てる方法も決定します。さらにバックエンドは、オペレーションの集合パターンが、最適化されたライブラリ呼び出しと一致するか確認します。

次に、ターゲット固有のコードを生成します。XLAに付随するCPUとGPUのバックエンドは、low-level IRとして LLVM を採用し、最適化やコード生成もLLVMを使って行います。これらのバックエンドは、XLAのHLO computationを表現するために効率的なLLVM IRを出力し、LLVMを使ってLLVM IRからネイティブコードを生成します。

現在、GPUのバックエンドはLLVM NVPTXのバックエンド経由でNVIDIA GPUをサポートし、CPUのバックエンドは複数のCPUのISAをサポートしています。

サポートするプラットフォーム

XLAは、x86-64とNVIDIA GPU向けに JITコンパイル機能 を、x86-64とARM向けに AOTコンパイル機能 をサポートしています。