Apache Software Foundation が管理する深層学習コンパイラスタック。モデルの計算グラフを入力として受け取り、CPU・GPU・NPU・FPGA 等の多様なハードウェア向けに自動最適化されたコードを生成する。AutoTVM と Meta Schedule による自動チューニングが特徴。
Apache TVM は 2018 年にワシントン大学の陳天奇(Tianqi Chen)らが開発した、深層学習モデル向けのオープンソースコンパイラスタックである。TensorFlow・PyTorch・ONNX 等のフレームワークからエクスポートされたモデルを入力とし、ターゲットハードウェアに特化した最適化コード を自動生成する。
TVM のコンパイルは以下のパイプラインで処理される。
| ステージ | 処理 |
|---|---|
| フロントエンド | TensorFlow/PyTorch/ONNX モデルを Relay IR に変換 |
| グラフ最適化 | 演算融合・定数畳み込み・レイアウト変換・量子化 |
| テンソル式生成 | 各演算をテンソル式(Tensor Expression)に分解 |
| スケジューリング | ループ分割・タイリング・ベクトル化・並列化のスケジュール決定 |
| コード生成 | ターゲットデバイス向けのネイティブコード出力 |
TVM は業界最多クラスのハードウェアバックエンドをサポートする。
この広範なハードウェアサポートにより、「一度書いたモデルをどこでも高速に動かす」というポータビリティを実現する。
TVM の最大の差別化要因は 自動チューニング 機能である。
MicroTVM は TVM のサブプロジェクトで、OS なしの組込みデバイス(マイコン:Cortex-M、ESP32 等)向けの推論コードを生成する。
ONNX Runtime は推論エンジンとして既存のバックエンド(cuDNN、TensorRT 等)を選択的に呼び出す方式。TVM は計算グラフを独自にコンパイルし、ターゲットハードウェア向けに最適化されたカスタムカーネルを生成する。TVM の方が最適化の自由度が高いが、セットアップとチューニングの手間がかかる。
モデルサイズとターゲットハードウェアに依存するが、ResNet-50 の全演算チューニングで 2-8 時間(GPU)、1-4 時間(CPU)が目安。Meta Schedule 導入後はチューニング時間が大幅に短縮された。チューニング結果はキャッシュ可能。
Apache License 2.0 で公開されており、商用利用に制限はない。OctoML(現 OctoAI)が TVM ベースの商用最適化サービスを提供していたほか、複数の半導体メーカーが自社チップ向けの TVM バックエンドを開発している。