こんにちは。NTT 研究所の森田です。先日、シアトルのワシントン大学で開催された TVM and Deep Learning Compiler Conference に招待されたので参加してきました。TVM って何？という方にもわかるように、TVM の概要から始めて、会議の内容までご紹介したいと思います。
深層学習コンパイラとは
TVM は深層学習コンパイラというジャンルに属するソフトウェアです（推論コンパイラ、深層学習オプティマイザなどと呼ばれることもあります）。「コンパイラ」と聞くと、ソースコードから実行プログラムへの変換、をイメージされるかもしれませんが、深層学習コンパイラの入力は、ソースコードではなく、深層学習の学習済みモデルになります。
深層学習によって AI を作るためには、大量のデータによって学習を行った学習済みモデルと、AI 処理を行うハードウェアが必要になります。Tensorflow などの学習フレームワークの普及により、学習済みモデルの入手は非常に容易になりました。一方、この学習済みモデルをハードウェア上でどのように処理させるか、という点に関しては未だに課題があります。学習フレームワークが動作するのは、通常は高性能な GPU が搭載された高スペック環境です。一方、AI 処理を行う環境はエッジ端末などの低スペック環境であることも多く、学習したモデルを効率的に動作させるためには、適切なモデルの変換や AI 処理プログラムの生成が必要になります。この作業を自動で行ってくれるのが深層学習コンパイラの役割です。
TVM とは
TVM はオープンソースで開発されている深層学習コンパイラです。同様のオープンソースソフトウェアはいくつかありますが、TVM は
- 対応する学習フレームワークやハードウェアの多様性に富んでいる
- 幅広い組織から開発者から参加しており、開発者数も多い
- 大学発のソフトウェアということもあり、最先端の研究成果が積極的に取り入れられている
といった特徴があります。
深層学習コンパイラに関しては、NVIDIA の TensorRT や Intel の OpenVINO のような、ハードウェアベンダが提供しているソフトウェアもあります。ただし、AI 処理を行う環境は様々で、最適なハードウェアの選択も状況に依存します。例えば NVIDIA の GPU は AI 分野においてデファクトスタンダードとも言えるハードウェアですが、場合によっては、より省電力で安価なハードウェアを利用したいこともあります。今後は、ハードウェア選択の中立性の観点から、ベンダに依存しない、オープンソースの推論コンパイラが主流になっていくと考えています。
NTT 研究所からは今春から TVM に取り組み始めまして、現在２名が TVM コミュニティでレビューアの役割を担当し、活動しております。
TVM の活用事例
深層学習の学習フレームワークは、GAFA に代表される大手クラウドベンダが、それぞれ別のソフトウェアを採用（Google の Tensorflow、など）していて、ソフトウェア間の競争が発生しました。実は深層学習コンパイラも同様の構図になっています。深層学習コンパイラにおいて、TVM を採用しているクラウドベンダは Amazon です。Amazon では様々な用途で TVM を活用しているという話が、会議の場でありました。
- Amazon の AI サービスの最適化のために利用（例：Amazon Echo でのリソース利用効率化、 Rekognition でのレイテンシ改善）
- Amazon Inferentia (AWS の AI チップ) の開発ツールとして利用
- Amazon SageMaker Neo (AWS のコンパイラサービス) での内部利用
Amazon 以外ですと、Huawei がかなり積極的に TVM を実用化しています。例えば Huawei の Atlas 製品シリーズでは内部で TVM を活用しているとのことです。また、Huawei は利用だけではなく、TVM の開発にも積極的で、TVM コミュニティに対してかなりのコミットを行っています。
AutoTVM
AutoTVM は TVM のコンパイル結果を自動チューニングする機能です。会議の場では、TensorRT や OpenVINO、Tensorflow XLA (Google が採用している深層学習コンパイラ) よりも TVM が高速であるといった、性能の優位性に関する発表が多数ありましたが、これらは AutoTVM 機能のおかげと言っても過言ではありません。従来は、より高速に AI 処理を行おうとするとどうしても全自動とはいかず、最終的には（職人のような方々の）人手によるチューニング作業が必要でした。AutoTVM はこのチューニングの箇所に機械学習を用いることで、自動で「最適なコード」を発見しようという機能です。この機能から TVM のことを「Learning-based Deep Learning Compiler」などと呼ぶ表現も会議の場でありました。
AutoTVM を利用するためには、利用者は scheduler template (何がチューニングパラメータかを記述したもの) を指定します。この scheduler template は通常は既存のものを利用しますが、必要であれば自分で記述することも可能です。AutoTVM は実測を通して性能のチューニングを行うので、この template があれば、新しいハードウェアにもそのまま対応することができます。また、scheduler template までも自動で生成するという auto scheduler という取り組みの紹介もありました。auto scheduler は近いうちにリリースされるということで楽しみです。
Relay
TVM の内部では、深層学習のニューラルネットワークを記述するための中間言語が実装されています。以前はこの中間言語は NNVM というものだったのですが、現在この NNVM を Relay という次世代の中間言語に置き換える作業が、TVM コミュニティ内で急ピッチで進められています。
この Relay ですが、「学習」を意識して設計されているというのが大きなポイントです。会議では TVM を学習に用いるロードマップの話がありました。また、TVM は様々なハードウェア向けにコンパイルできるので、TVM を学習に利用できるとなると、FPGA や ASIC を用いた学習も行えるようになります。ここまでいくと、Google の TPU 環境のような機械学習環境が、オープンソースソフトウェアと専用チップで容易に構築できる時代になりそうです。
また、Relay は量子化機能（バイナリネットに関する LT 発表が 3 件もありました）やヘテロデバイス構成での実行サポートも設計に取り入れようとしており、これらの機能は今後の TVM での開発の中心になりそうです。
VTA
VTA と書いて「ヴィータ」と読みます。VTA は TVM でハードウェアアクセラレータを扱うためのフレームワークです。現在は Xilinx 社の FPGA を用いた PYNQ ボードをターゲットとし、FPGA 上に作成した行列乗算器等のアクセラレータ（高位合成で作成）を利用して深層学習を高速化する実装が用意されています。
この VTA を用いて、
- Huawei Da Vinci
- Qualcomm Hexagon
- Amazon: Inferentia
などの AI チップから TVM を利用することが可能になり、その取り組みの紹介が会議でありました。VTA が今後の AI チップの深層学習コンパイラとして、デファクトスタンダードになりつつある印象を受けます。
また、ハードウェア生成そのものまでも AutoTVM と同じように機械学習で行ってしまうという AutoVTA という話もありました。AI ハードウェアの回路構成をニューラルネットワークに応じて作り変えるということで、とても興味深いです。
現在 VTA が対応しているハードウェアは前述の PYNQ だけですが、現在、Intel の FPGA を用いた DE10-NANO への対応も進められています。近日中にリリースされると言われている TVM v0.5 には、この DE10-NANO 対応も取り入れられる予定とのことです。
高位合成バックエンド
会議では私からも「OpenCL backend for FPGA」というタイトルで発表を行いました。
現在の TVM で FPGA を扱おうとすると、上述の VTA を利用するのが基本的なアプローチなのですが、私が発表したのは Xilinx や Intel が提供している高位合成ツールを用いてコンパイルを行おうというアプローチです。これにより、FPGA も GPU や CPU と同じように扱えるようになるのでソースコードの共有化もできますし、高位合成環境に含まれるプロファイラなど強力な開発ツールも利用することが可能になります。
高位合成を利用することの大きな懸念は性能やリソース使用量ですが、手動で最適化することでかなりいい線までいけそうなことが見えています。今後、これをどうやって TVM 上で自動化していくかが課題です。
発表後に TVM のコミッタ達と議論して、スタンフォード大学が開発している高位合成言語である Spatial もおすすめされました。Spatial は会議で発表もあったのですが、高位合成のために一から設計された言語で、我々としても注目しています。
おわりに
TVM コミュニティのコミッタ達は、とても優秀なのはソースコードを見てわかるのですが、会って話してみると、人格者でとてもいい人たちでした。こういった、人の魅力もこのコミュニティの魅力のひとつですね。今回の TVM Conference は第一回でしたが、第二回も一年後あたりに行う予定とのことです。
海外では大きな企業による実用化事例も出てきている TVM ですが、国内でも盛り上げていけたらと思っています。もし興味を持たれた方がいらっしゃいましたら、是非ご連絡ください。一緒に仕事をしてみたいという方は弊社の採用情報ページも是非！
