[論文レビュー] Tensor Comprehensions: Framework-Agnostic High-Performance Machine Learning Abstractions

Nicolas Vasilache, Oleksandr Zinenko+7 more|arXiv (Cornell University)|Feb 13, 2018

Computational Physics and Python Applications参考文献 61被引用数 252

ひとこと要約

Tensor Comprehensions (TC) を導入し、ML ワークロードの高性能 CUDA カーネルを生成するドメイン特化言語およびポリヘドラル JIT コンパイラで、クロスフレームワーク統合と演算子融合を実現します。

ABSTRACT

Deep learning models with convolutional and recurrent networks are now ubiquitous and analyze massive amounts of audio, image, video, text and graph data, with applications in automatic translation, speech-to-text, scene understanding, ranking user preferences, ad placement, etc. Competing frameworks for building these networks such as TensorFlow, Chainer, CNTK, Torch/PyTorch, Caffe1/2, MXNet and Theano, explore different tradeoffs between usability and expressiveness, research or production orientation and supported hardware. They operate on a DAG of computational operators, wrapping high-performance libraries such as CUDNN (for NVIDIA GPUs) or NNPACK (for various CPUs), and automate memory allocation, synchronization, distribution. Custom operators are needed where the computation does not fit existing high-performance library calls, usually at a high engineering cost. This is frequently required when new operators are invented by researchers: such operators suffer a severe performance penalty, which limits the pace of innovation. Furthermore, even if there is an existing runtime call these frameworks can use, it often doesn't offer optimal performance for a user's particular network architecture and dataset, missing optimizations between operators as well as optimizations that can be done knowing the size and shape of data. Our contributions include (1) a language close to the mathematics of deep learning called Tensor Comprehensions, (2) a polyhedral Just-In-Time compiler to convert a mathematical description of a deep learning DAG into a CUDA kernel with delegated memory management and synchronization, also providing optimizations such as operator fusion and specialization for specific sizes, (3) a compilation cache populated by an autotuner. [Abstract cutoff]

研究の動機と目的

• ML のためのテンソル演算を表現する、数式に適した簡潔な言語を提供する。
• テンソルコンプリヘンションを、メモリ管理とスケジューリングを備えた最適化された GPU コードへ落とす。
• TCを PyTorch および Caffe2 と連携させることで、クロスフレームワーク統合を実現する。
• 最適化機会を探索するオートチューナを介した自動チューニングを提供する。
• 既存の ML フレームワークとの競争力のある性能と実用的な統合を示す。

提案手法

• テンソル演算のための、アインシュタインの総和に近い表記である Tensor Comprehensions (TC) を定義する。
• メモリ管理と同期を伴う TC を CUDA カーネルへ低減するポリヘドラル JIT コンパイラを開発する。
• カーネル融合とサイズ固有の特殊化を備えた、特化したポリヘドラル最適化フローを作成する。
• JIT コンパイルとコンパイルキャッシュを活用したオートチューニングフレームワークを実装する。
• in-process インターフェースと ATen 非同期テンソルライブラリを介して、PyTorch と Caffe2 と TC を統合する。

実験結果

リサーチクエスチョン

• RQ1 TC は共通の ML 演算子やカスタム演算子を簡潔かつ安全に表現できるか？
• RQ2 TC ベースのフローは、標準的な ML カーネルに対して、GPU 上でベンダーライブラリと競合する性能を達成できるか？
• RQ3 ポリヘドラルベースのオプティマイザは、深層学習モデルにおけるデータサイズ/形状に対して効果的に演算子を融合し、最適化できるか？
• RQ4 TC は PyTorch や Caffe2 などの主流フレームワークとどれくらいうまく統合され、実運用志向の道筋を提供できるか？
• RQ5 自動チューニングと JIT コンパイルが、ハードウェア固有の性能向上を達成するうえでどの程度影響を与えるか？

主な発見

• TC フローは ML ワークロードに関連するカーネルで NVIDIA ライブラリより最大 4x のスピードアップを達成する。
• Caffe2 と PyTorch とのエンドツーエンドの TC 統合は、実用的で生産志向の適用性を示す。
• 長依存のディープラーニングカーネルに対して効果的な最適化を提供する、特殊化されたポリヘドラルコンパイラ。
• オートチューニングとコードキャッシュは、非標準のサイズやレイアウトに対するドメイン特化の最適化を実現する。
• フレームワークは安全なインプレース更新と、共通レイヤー（例: SGEMM, conv2d, maxpool）の単純で宣言的な定義をサポートします。
• システムはコア部でフレームワーク非依存性を保ちながら、実運用環境との緊密な統合を提供します。