このようにしてまとめた実装では、不要なメモリ割り当てを行わずに、1 個にまとめられた要素ごとのループ内で最終結果を計算できます。さらに高度なシナリオでは、このような演算が行列の乗算にまとめられることもあります。
チーム内でベンチマーク テストを行ったところ、XLA は従来の TensorFlow に比べて NVIDIA GPU 上で最大 50% のスピードアップを実現できました。事前の予想どおり、もっとも大きな高速化が達成されたのは、要素ごとの演算が長く続くモデルです。このようなモデルはより効率的なループにまとめられます。ただし XLA はまだ試験運用版であり、ベンチマークによってはスピードが低下する可能性もあります。
Chris Leary と Todd Wang による TensorFlow Developer Summit の講演では、TensorFlow で XLA、JIT、AOT などのコンパイル技術を活用し、実行時間を最低限に抑えつつ計算リソースを最大限に活用する方法が説明されています。
徹底した特化処理による実行サイズの削減
XLA は、パフォーマンスの改善だけでなく、実行サイズの削減機能も提供しているため、限られたメモリしか搭載されていない環境（モバイル端末など）向けの TensorFlow モデルにも有用です。
tfcompile は、XLA を活用して事前コンパイル（AOT）を行うためのツールです。グラフ全体が XLA にコンパイルされた後、グラフ内の演算を実装するコンパクトなバイナリコードが生成されます。ランタイムを最小限にとどめる機能をこのスキームに組み合わせることで、バイナリのサイズを大幅に削減できます。
たとえば、android-arm で 3 層、60 ノード幅の LSTM スタック モデルを作成したところ、オリジナルの TF モデルのサイズは 2.6 MB（ランタイム 1 MB + グラフ 1.6 MB）でしたが、これを XLA でコンパイルすると 600 KB まで縮小されました。
このサイズ削減は、静的コンパイルによるモデルの完全な特化によって実現されています。個々のモデルを実行する際には、TensorFlow ランタイムの処理能力と柔軟性をフルに維持する必要はありません。ユーザーが実際のグラフの実装で使いたい演算のみがバイナリコードにコンパイルされるからです。もっとも、XLA の CPU バックエンドが生成するコードのパフォーマンスはまだまだベストとは言えず、今後も引き続き改善していく予定です。
異なるバックエンドや装置のサポート
TensorFlow グラフを新たな種類の計算デバイスで実行するには、新しいデバイス用にすべての TensorFlow 演算（カーネル）を再実装する必要があります。デバイスによっては、これに膨大な作業が必要となります。
XLA では、カスタム バックエンドを追加することで簡単に新しいデバイスをサポートできます。TensorFlow は XLA をターゲットにできるので、新しいデバイスのバックエンドを XLA に追加すれば、そのデバイスで TensorFlow グラフを実行できます。前述のとおり、XLA は複雑な演算をプリミティブ演算に分解するため、
XLA の演算はプリミティブ演算のみとなります。そのため、新しいデバイスのサポートに必要な作業ははるかに少なくなります。XLA にカスタム バックエンドを追加するプロセスの詳細は、
こちらのページに記載されています。ちなみに Google 社内では、
TPU をターゲットにするため、XLA のこのメカニズムを活用しています。
まとめと今後の展望
まだ開発の初期段階の XLA ですが、ユースケースによっては非常に大きな効果を発揮でき、TensorFlow が今後この技術から大きなメリットを得られることは明らかです。そこでGoogle では、XLA を TensorFlow の GitHub リポジトリ で
先行リリースすることを決定しました。今後は開発者コミュニティからの開発支援も受けつつ、TensorFlow をさまざまな計算デバイス向けに最適化し、ランタイムとモデルを新しい種類のハードウェアで実行するための便利な基盤を提供していきたいと考えています。
Reviewed by
Kaz Sato - Staff Developer Advocate, Google Cloud
