[論文レビュー] DLVM: A modern compiler infrastructure for deep learning systems

Richard Wei, Lane Schwartz+1 more|arXiv (Cornell University)|Nov 8, 2017

Parallel Computing and Optimization Techniques参考文献 4被引用数 34

ひとこと要約

DLVM は、深層学習向けに設計された現代的で LLVM を模倣したコンパイラインフラで、テンソルDSLの型安全でモジュラーかつ最適化可能な基盤を提供する。線形代数の中間表現、随伴コード生成による原理的で整合性の取れた自動微分、および LLVM を介した GPU ターゲット向けコード生成を通じて、高パフォーマンスで安全かつ合成可能な深層学習システムを実現する。Swift で実装されたプロトタイプDSL（NNKit）により、エンド・ツー・エンドの最適化と JIT コンパイルが実証された。

ABSTRACT

Deep learning software demands reliability and performance. However, many of the existing deep learning frameworks are software libraries that act as an unsafe DSL in Python and a computation graph interpreter. We present DLVM, a design and implementation of a compiler infrastructure with a linear algebra intermediate representation, algorithmic differentiation by adjoint code generation, domain-specific optimizations and a code generator targeting GPU via LLVM. Designed as a modern compiler infrastructure inspired by LLVM, DLVM is more modular and more generic than existing deep learning compiler frameworks, and supports tensor DSLs with high expressivity. With our prototypical staged DSL embedded in Swift, we argue that the DLVM system enables a form of modular, safe and performant frameworks for deep learning.

研究の動機と目的

• 既存の深層学習フレームワークの限界を是正すること。これらはしばしば不安全で動的型付けのDSLをPythonで使用しており、モジュラーかつ合成可能なコンパイラ設計を欠いている。
• 完全な機能を備えた静的型付けのDSLをサポートする、原理的でモジュラーかつ拡張可能なコンパイラインフラを提供すること。
• カーネル統合、代数的単純化、ADチェックポイント技術を含む高度な最適化を通じて、高パフォーマンス実行を実現すること。
• 自動微分をコンパイラパイプライン内での一等級の変換フェーズとして統合し、正しさとパフォーマンスの両方を向上させること。
• LLVMとの統合を通じて、CPU や GPU を含む多様なハードウェアターゲット向けに、ポータブルで効率的なコード生成を可能にすること。

提案手法

• 制御フローを備えた SSA 形式のテンソル計算用ドメイン固有中間表現（IR）の設計と実装。これにより、標準的なコンパイラ最適化が可能になる。
• 自動微分のためのソースコード変換。演算子オーバーロードに依存せず、随伴コードを第一級のコンパイラフェーズとして生成する。
• LLVM との統合による低レベルコード生成。CPU や GPU（NVIDIA GPU は NVPTX を介して）への透明なターゲティングを可能にする。
• Swift で実装された段階的DSL（NNKit）の開発。Rep<T>ラッパーを用いて、テンソル計算の静的解析、型安全性、遅延評価を実現する。
• 4段階のJITコンパイルパイプライン：式の段階的生成、形状の特殊化、DLVM IR への低レベル化、関数の再結合化によるSwiftでのコンパイル関数の公開。
• インライン化、ループ最適化、計算カーネル統合といった一般およびドメイン特化最適化を、モジュラーかつ拡張可能なコンパイラフレームワーク内で適用する。

実験結果

リサーチクエスチョン

• RQ1 確立されたコンパイラ手法に基づく現代的でモジュラーなコンパイラインフラは、深層学習フレームワークの信頼性とパフォーマンスを顕著に向上させ得るか？
• RQ2 自動微分を、ライブラリレベルの機能ではなく、深層学習コンパイラ内での第一級の変換フェーズとして統合するにはどうすればよいか？
• RQ3 システム言語（例：Swift）で実装された静的型付けで段階的なDSLは、深層学習において安全なプロトタイピングと高パフォーマンス実行を両立できるか？
• RQ4 統一されたIRと最適化パイプラインは、多様なハードウェアターゲット上で、高レベルのニューラルネットワーク計算と低レベルのカーネル生成の両方を効果的にサポートできるか？
• RQ5 IRにSSA形式と制御フローグラフを採用することで、グラフベースの表現と比較して、深層学習IRの表現力と最適化可能性はどのように向上するか？

主な発見

• DLVM は、LLVM を模倣した原理的で洗練されたコンパイラインフラを通じて、モジュラーで安全かつ高パフォーマンスな深層学習フレームワークの実現に成功した。
• 自動微分をライブラリ機能ではなくコンパイラフェーズとして統合することで、勾配計算の最適化をより積極的かつ正確に実現できるようになった。
• Swift で実装された段階的DSL（NNKit）は、静的解析と型安全性を完全に実現しながら、DLVM を介した効率的なJITコンパイルとコード生成を可能にした。
• DLVMのIRは、カーネル統合、代数的単純化、ADチェックポイント技術といった高度な最適化をサポートしており、カーネル起動のオーバーヘッドを低減し、パフォーマンスを向上させた。
• LLVM を用いた高レベルのテンソル式から最適化されたGPUカーネルへのエンド・ツー・エンドコンパイルが実証され、NVIDIA GPU への対応に加え、HPVM などの他のバックエンドへの拡張性も示された。
• プロトタイプ実装により、コンパイラベースのアプローチが深層学習システムの開発を統合的かつ簡素化でき、既存フレームワークの複雑さと脆弱性を軽減できることが示された。