TVMとは
ペーパーはこちらからダウンロードできます。 https://arxiv.org/abs/1802.04799
TVMは、ASICやAPGAの様な組込様CPU等、様々なバックエンドでDeep Learningを動かすことを目的としたコンパイラです。 グラフレベルや演算子レベルの様々な最適化を行います。
興味本位でペーパーを読んでみたら、僕の知らない最適化テクニックが使われていたので紹介します。 他にも面白いことは書いてあったのですが、最も気になった2つを紹介します。
TVM stack
TVMは図のようなスタック構造になっています。上から順番に最適化をしていって、最後にLLVMやCUDAのソースコードを生成します。High level Graph Rewritingの話(Section 3)や、Tensorとスケジューリングの話(Section 4)で最適化した後、それらの組み合わせを考慮するときに、機械学習を使って最適化をしています。
TVMは、このOptimizerに来る前に可能な限りの最適化を考えます。その結果、10億以上の最適化の組み合わせ(configuration)から、実際のワークロードに適した組み合わせの検索をOptimizerが行います。つまり 、Optimizerは実際に最適化をするのではなく、最適化の組み合わせの探索に機械学習を使っています。
ML based Cost Model
最適化の探索には、様々なファクターがあります。 - メモリアクセスパターン - データの再利用 - パイプラインの依存性 - スレッドのパターン - etc 最適化のするに当たっては計算のコストモデルを作る必要があります。 従来のAuto-turningの手法では良い組み合わせを探すのに多くの実験が必要になります。定義済みのコストモデルを使う場合は、バックエンドのハードウェアへの依存性が高くなります。そこでTVMでは、機械学習を使ったアプローチを取っています。
図のように、ループの構文ツリーをTreeRNNに入れ、特徴量を出しコストを予想しています。
Schedule Exploration
コストモデルが決まった後は、そのコストモデルに従って一番良い組み合わせを探します。理想的には全てのコンフィグレーションを試し上位のいくつかを選択すれば良いのですが、それには強大な探索空間が必要になります。その問題を解決するために、TVMは a parallel simulated annealing algorithmを使います。
まずはランダムに選んだコンフィグレーションから開始し、ランダムウォークをしながら、コストモデルが予測するコストが低い方へ移動していきます。ランダムウォークはコストの低いところを見つけ、移動しなくなるまで続けられます。
感想
コンパイラの最適化技術がここまで進んでいるとは思いませんでした。個別のターゲットに対して、個別の最適化を行う時代は過ぎているだと思います。AIによって世界中の誰でも、どのデバイスでも、最適化職人の技術の恩恵にあずかれると思うと素晴らしいですね。
コンパイラ勉強会が来月に開催されるようで、申し込んではみましたが抽選に通るかどうか難しいそうです。 https://connpass.com/event/103976/
