TVMの次期自動チューニング機構（Ansor）について
皆さんこんにちは、NTTの森永です。私の記事は前回、前々回に引き続き、今回も深層学習コンパイラOSSのTVMについてです。
TVMでは、AutoTVMという独自の仕組みにより、畳み込みや行列積などの演算の最適化を自動で行います。AutoTVMは機械学習の推論の最適化のために機械学習を用いており、すでにこの時点で技術的に興味深いのですが、6/18にAutoTVMの提案者からAnsor（AutoTVM v2.0）という新技術についてのTVMコミュニティへの投稿とarxivへの技術論文掲載が行われました。それらによると、Ansorは、AutoTVMよりも高速な最適化をより短いチューニング時間で達成でき、更にAutoTVMに必要だった最適化のためのテンプレートが必要なくなるとのことです。
既にTVMコミュニティの間では、TVMにAnsorを導入していくことで合意されています。Ansorのコードは分量が多く何回かにパッチを分けてマージされていく計画のため、実際に使えるようになるのは少し先でしょう。ですが、長い間TVMの高速化の核でもあったAutoTVMに取って代わるとなると、深層学習コンパイラや推論高速化の分野における今後の要注目技術であることは間違いありません。一足先にキャッチアップするため、本記事では、（論文の翻訳や詳細な解説というよりは）AnsorがAutoTVMと比べどう優れるのか？TVMがどう変わっていくのか？という観点で論文から、Ansorの技術のポイントを紹介しようと思います。
なお、RFCでは括弧内でAutoTVM v2.0とも称されていましたが、実際の仕組みは大きく異なっており、Ansorにも直接関わっているコントリビュータから、v2.0という表記は適切でないという指摘もありました。
I think framing Ansor as AutoTVM v2.0 is somewhat misleading as Ansor is taking a totally different approach from AutoTVM.
1. AutoTVMの問題点
AutoTVMは以下のような問題を抱えています。
問題1：AutoTVMテンプレートの作成が難しい
AutoTVMでは、畳み込みなどの各演算に対して、ハードウェアごとに最適化のためのテンプレートが必要です。このテンプレートでは、キャッシュの利用、ループタイリングの指定などのプログラムの最適化と、チューニングの探索空間（タイリングの大きさの候補等）を手動で定義します。そのため、アルゴリズム、各ハードウェアごとのプログラム最適化のコツ、TVMの記法全てに詳しい人間による実装が必須で、テンプレートの作成はかなり難易度が高いです。そして、新しいオペレータが増えれば、ハードウェアの数だけ新しいテンプレートの作成が必要となり、多大な実装労力がかかります。
問題2：探索空間の範囲が狭く、優れた解が得られないことがある
AutoTVMやFlexTensorのような、テンプレートを用いるチューニング方式ではテンプレートの作成者が設定したプログラム候補しか探索空間に入らず、探索範囲が限定的です。より高性能な解（最適化後のプログラム）を見つけるには、大きな探索空間上での探索が必要かもしれません。
問題3：チューニングが遅い
AutoTVMでは、例えば、１つのネットワークの最適化に丸1日あるいはそれ以上のチューニング時間がかかってしまうことが少なくありません。これは、探索中の候補に対して毎回コンパイルと複数回の実測が必要で、それを何百あるいは何千ステップも行うためです。
Ansorではこれらの問題が一気に解決します。問題1 問題2については、独自の階層構造により膨大な探索空間を表現し、かつそれをテンプレートなしで自動で生成します。天下り的ではありますが、Ansorの新しい探索空間は従来のテンプレートの探索空間に含まれない高速な解を含むことが結果から確認されています。問題3については、膨大な探索空間を進化的アルゴリズムを用いて効率よく探索し、実測回数を減らす工夫によりチューニング時間を抑えています。さらに、AutoTVMにはなかったタスクスケジューラの仕組みを考案し、実時間を効率よく使ってチューニングの効果を最大化します。
では、これらの「探索空間」「探索アルゴリズム」「タスクスケジューラ」の特徴を見ていきましょう。
2. 探索空間
Ansorの探索空間は、計算の定義や計算グラフを元にSketchとAnnotationという二つの階層で構成されます。Sketchは高レベルな構造を表し、Annotationは細かい最適化の要素を表すものです。ある計算に対するSketchとAnnotationの実際のイメージは、論文中の図5を見ると、掴み安いと思います。
Sketchは、計算グラフに対して、各ノードを計算するための指針を大まかにと定めるものです。Sketchは、対象の計算グラフからルールベースで自動生成されます。Sketchの生成は、計算グラフの出力ノードから順に再帰的に、各ノードに対する適用可能なルールを全て調べ、全てのノードに対するルール適用を行うことで、候補をチューニングの前にあらかじめ全列挙します。この生成ルールはハードウェアごとに定義されます。例えば、CPUにおける生成ルールには以下のようなものが含まれています。
Skip : 特にルールを適用せず、次のノードへ行きます。
Always Inline : そのノードの計算を全てインライン展開します。要素ごとの和やReluなど、要素ごとに処理できる計算ノードに適用できます。適用後、次のノードへ行きます。
Multi-level Tiling ：そのノードの計算の全てのループをタイリングします。行列積やconv2dのようなデータ再利用の機会が多い処理の計算ノードに適用できます。適用後、次のノードへ行きます。
このように、Sketchの生成ルールは、あるノードを計算する上で「インライン展開する」「タイリングをする」「タイリングし、直後のノードとfuseする」など、ざっくりと有効そうな最適化の方針を定めるものとなっています。
さて、ここまでのSketchだけでは、具体的なタイルサイズの情報や、ベクトル化や並列化等のアノテーション指定など最適化情報は含まれません。Annotationは、Sketchに付与され、そのような最適化の詳細を示すものです。Sketchと比べ、Annotationの候補は全て調べ上げ、保持できるような数ではありません。そのため、後に解説するAnsorの探索ではランダムに選んだスケッチに対してランダムにAnnotationを生成する、という方法で複数のプログラムサンプルを生成し、それらをベースに微調整する、という方針が取られます。Annotationの生成では、全列挙された中からランダムに選んだ Sketchに対して、タイルサイズ等をランダムに埋めたのち、外側のループの並列化、内側のループのベクトル化、内側ループの数回のアンロールのアノテーションをランダムに適用します。また、ループ回数が1と定まったタイルについては省略を行います。
AutoTVMと比較して見ると、各ハードウェアごとで有効な著名な最適化手法についてはSketchの生成ルール内に選択的にノウハウを反映しつつ、各テンプレートの内容が担っていた細かいチューニングの仕方と探索空間の表現はAnnotationが担っていると言えます。例えば、Annotationで指定されるタイルサイズはAutoTVMではdefine_knobで定義されていた探索空間そのものですし、ベクトル化やアンロールの指定は、define_annotateというAPIを用いてテンプレート中で指定されていたものです。我々が普段、手動で計算プログラムを最適化をする際も、全体の方針はよく知られた最適化の手法から選び、細かい部分については試行錯誤の中でノウハウを獲得していくことがあるでしょう。AnsorにおけるSketchとAnnotationの概念はこれに相応していると言えそうです。
３. 探索アルゴリズム
Ansorが膨大な広さの探索空間を用いる分、チューニングの難易度は大幅に上がるはずです。その上でAnsorがAutoTVMよりも優れた解を短時間で発見する仕組みの鍵は、独自の進化的アルゴリズムにあります。進化的アルゴリズムは、詳しくは省略しますが、生物の進化の仕組みを模倣した探索アルゴリズムの総称です。ざっくりとイメージを掴みたい人は、以下に紹介する動画を見るのがオススメです。
進化的アルゴリズムでは、各世代のサンプルたちを評価し、優れたものが残りやすいよう選別します。また、選別された次世代にランダムな変化（進化）や新しいランダムサンプルを加えていきます。Ansorでは、この際の変化に、以下のように交配（crossover）や突然変異（mutation）といった進化的アルゴリズムの考え方が含まれているのがポイントです。
Tile Size Mutation：タイル化されたループをランダムに選び、その片方のループ長を定数で割り、もう片方を同じ定数倍します。
Parallel-Vectorize Mutation：ランダムに選んだループに、ベクトル化や並列化のAnnotationをつけます。
Node-base Crossover：Ansorでは、各サンプルのSketchおよびAnnotationの書き換えの変更過程が保持されています。この過程を遺伝子とみなして、2つのサンプルの書き換えられた過程をランダムに組み合わせします。
これらの突然変異や交配を繰り返しながら、選択により何世代も解候補を洗練させていくことで、Ansorでは、効率よく優れた解が得られます。
ところで、AutoTVMの論文では、進化的アルゴリズム（の一種のGA）よりも焼きなまし法での探索の方が有効性が高い結果が報告されていました。進化的アルゴリズムは、解空間の構造が複雑だがなんらかの秩序があり、かつ探索空間が膨大な探索問題に対して有効なことが多いと言われています。従来のAutoTVMでは、最適化のノウハウは主にテンプレートに反映されており、探索空間にはあまり含まれていません。これに対し、Ansorでは、最適化のノウハウを細かく反映する必要のある解空間の複雑さと、解空間自体の膨大な広さを持つため、進化的アルゴリズムが効果的に機能したのだと考えられそうです。
また、Ansorの進化的アルゴリズムの繰り返しでは、数世代分の過程での世代の評価と進化は全てコストモデル上で行い、数世代に一度の世代のみ実測に基づいてコストモデルの更新を行う方式を採用しています。コストモデルでの評価は、コンパイルを伴うハードウェアでの実測よりも早いので、これにより実質的には大量の候補の評価を行いつつも、大幅にチューニング時間を削減することができます。
このようにAnsorでは、進化的アルゴリズムを採用し、独自の進化の仕組みの設計とコストモデルを用いて実測回数を抑える工夫により、広い探索空間でも、高性能かつ高速なチューニングを可能にしています。
4. タスクスケジューラ
AutoTVMでは、すぐにチューニングが収束するレイヤーも収束に長い時間がかかるレイヤーもあまり区別せず、同じステップ数の探索を行う使われ方が主流でした。また、現実的には、チューニングによる実行時間の改善が大きいレイヤーほど集中的にチューニングしたいはずですが、AutoTVMにはその仕組みはありません。Ansorでは、実時間を効率よく使って、最大限にチューニングするため、タスクスケジューラが導入されます。
Ansorのタスクスケジューラでは、チューニング全体を目的関数を最小化する時間配分問題として定式化します。目的関数の設定次第で、例えば、DNN全体の実行時間を一定時間で最小化したり、あるいは遅延基準を満たす範囲で短いチューニングを目指したり、参考値に対する幾何平均比で高速化率を最大化したり、と目的に応じた効率の良いチューニングが可能になります。
タスクスケジューラの実際の仕組みは、最初は全タスクに均等に単位時間を配分して、以降は各タスクに対して「次に時間を割り当てると目的関数がどのくらい減るか」の推測値を算出し、最も良さそうなタスクに高い確率で次の単位時間を割り当てることを繰り返します。この際の推測値の算出は、「そのタスクが早くなることで目的関数がどれぐらい減るか」「そのタスクが直近の時間割り当てでどのくらい早くなったか」「そのタスクの計算量（FLOPS）はどの程度か」「類似タスクがチューニングでどの程度早くなったか」といった情報から近似式を用いて計算されます。時間割り当ての確率の制御にはε-greedy法が用いられるようです。
実際の挙動では、例えば、目的関数をDNN全体の実行時間とすると、スケジューラは序盤は実行時間の大きいレイヤーに優先的に時間を割り当てます。そのレイヤーたちのチューニングによる改善が減ってくると、他のレイヤーにチューニング対象が移る、という動作がされます。これにより、期待値の高いレイヤーから優先的にチューニングされていくわけです。
5. Ansorの性能
論文中では、Intel CPU 、NVIDIA GPU、ARM CPUにおいて、最適化ライブラリ（TensorRT、cuDNNやMKL-DNNを用いるPytorch/Tensorflow）およびAutoTVMやFlexTensor、Halide auto-schedulerといった先行研究たちと性能を比較しています。いずれもfloat32の精度での評価となっています。
単独のオペレータの評価では、10種類のオペレータと２種類のバッチサイズ（1, 16）を組み合わせた20種類のケースをIntel CPUで評価し、19種のケースにてAnsorが最速（他フレームワークの1.1~32倍程度）であることが確認されています。多くのケースでAnsorの見つけた解は他の先行研究の探索空間の外にあり、広い探索空間が性能に寄与していることを示唆しているようです。
ネットワーク全体では、Intel CPU / NVIDIA GPU / ARM CPUにおいて、ResNet-50やDCGAN、BERTなど5種類のネットワークを２種類のバッチサイズで評価したところ、24/25のケースで他フレームワークやAutoTVMと比べ1.0~9.4倍速い性能が達成されています。最も優れた他の手段と比べると、幾何平均で、Intel CPUでは3.8倍、NVIDIA GPUでは 1.7倍、ARM CPUでは2.6倍の高速化が達成されています。
チューニング中の試行回数については、AutoTVMが50000回近くの測定で発見したmobilenet-v2やResNet-50におけるチューニング結果の性能を、1/10以下の測定回数で達成したことが述べられています。AutoTVMでは、DNN全体となると日数単位でチューニングが必要なこともあったので、それらが数時間で終わるとしたら、クラウドの様々なハードウェアを使って比べる、みたいなことも料金を気にせず従来より気軽に試せそうで楽しみです。
5. 終わりに
本記事では、Ansorの3つの特徴である、SketchとAnnotationの階層構造を持つ高い表現能力の探索空間、実測の回数を抑えた進化的アルゴリズムによる効率の良い探索、目的に応じて実時間を効率よく使えるタスクスケジューラをざっくり紹介しました。Ansorの詳細については、興味を持ったらぜひ元論文を細かく読んでみることをオススメします！
Ansorの仕組みは紹介したようにAutoTVMと大きく異なっていましたが、AutoTVMとほぼ同じ手順とAPI（タスク抽出→各タスクを最適化）で用いれられるそうです。AutoTVMはこれまで何度も触ってきたので、ユーザー目線ではこれは非常にありがたいですね。
また、Ansorの内容を踏まえると、単純に従来よりもTVMの性能が高くなるだけでなく、膨大な数の最適化テンプレートが必要なくなり、新しいオペレータやハードウェアの登場への適応サイクルが早くことも期待できそうです。あるいは、チューニングの恩恵に預かるまでに必要な実装量が大きく減ることから、独自アクセラレータのためのコンパイラとしてTVMの選択肢が有力になるかもしれません。
NTTでは、共にソフトウェアの研究やOSSコミュニティ活動を行う仲間を募集しています。興味のある方は、新卒中途問わず、ぜひ採用情報をチェックしてください！
