TensorFlow XLA について.

言わずと知れたTensorFlowというDeepLearningのOSSライブラリに、TensorFlow用のコンパイラであるTensorFlow XLAというものがある. 個人的にはDeepLearningや機械学習というものよりもコンパイラ周りの低レイヤのプログラミングに興味があるのでTensorFlowに最適化されたコンパイラはとても興味深い. DeepLearningは個人的な感覚として致命的に処理が遅いことと出力されるファイルサイズが巨大であることが扱いづらさのポイントであった. XLAはそれらの問題を解決するために特化したコンパイラであると言えると思う.

TensorFlow XLAは一言で言うと、Linear Algebraのドメイン特化したコンパイラフレームワークである. 通常のTensorFlowは、TensorFlow Runtime上で動いているため、 TensorFlow Runtimeのオーバーヘッドがあり、冒頭で書いた問題が存在してしまう. しかしXLAを導入することによって、XLAがTensorFlowコードをCPUの実行バイナリに変換するのでTensorFlow Runtimeを使わなくて済むようになり、TensorFlow Runtimeのオーバーヘッドを無くすことが可能となる. また、CPUでの実行バイナリとなるため、様々なデバイスのCPUコアでもTensorFlowのコードが動くことができる.

XLAはバックエンドとしてLLVMを利用している. TensorFlowのモデルをインプットとして受取り、XLAの中間表現に変換した後、最適化処理を行う. この最適化のあとで、LLVM IRに変換し、LLVMのpassに渡されて最終的にバイナリコードで出力される、という流れ. LLVMを利用して、x86, x86_64, ARM, ARM64, GPUなどの様々なアーキテクチャに対応することができる. ポイントとしては、LLVMレイヤ上でLinear Algebra向けの中間表現を実装しているという点で、これによって性能向上を大幅に上げることができている. また、XLAはJITとAOTの両方をサポートしており、 JITでは実行時のパラメータを利用した最適化による学習フェーズの高速化が可能となり、AOTではTensorFlow Runtimeから独立したバイナリファイルを生成することができるので、推論フェーズの低レイテンシ・省メモリ化が実現できるようになる.

このXLAを利用することで、NVIDIA GPU上の従来のTensorFlowのベンチマークと比べて最大50%の性能向上が見られたらしい. (ただし、ケースによっては性能が落ちる場合もあったとか.)

スマートフォンなどの軽量デバイス上でDeepLearningを動かそうとした場合、この分野はやっぱり大事になってくるだろう(IoT絡みでも). その際のドメイン特化のコンパイラというものはとても興味深い. LLVMにバックエンドを任せるのは今の所ベストな解だろうな、と思うのと、フロントエンドでどのような最適化を行うのか、というところが気になる.

そういえば、NVIDIA GPU向けのLLVM NVPTXをサポートしているようだけれど、この辺のLLVMの情報は全然得られていないなぁ、触らないと. やっぱり機械学習そのものよりも低レイヤの技術の方が断然面白くて楽しいですね.

One clap, two clap, three clap, forty?

By clapping more or less, you can signal to us which stories really stand out.