[論文レビュー] Tenspiler: A Verified Lifting-Based Compiler for Tensor Operations (Extended Version)

Jie Qiu, C.R. Cai+4 more|arXiv (Cornell University)|Jan 1, 2024

Parallel Computing and Optimization Techniques被引用数 3

ひとこと要約

Tenspiler は、Tensir と呼ばれる新規で統合的な中間言語を用いて、順次的な C++ や Python コードを最適化されたテンソル演算に変換する検証済みのリフトベースのコンパイラである。コード変換を構文誘導型プログラム合成問題として定式化し、正しさを保証する定理証明を活用することで、6 つのソフトウェアおよびハードウェアバックエンドで 69 個の実世界ベンチマークに対して平均 105× のカーネル実行速度向上と 9.65× のエンドツーエンド実行速度向上を達成した。

ABSTRACT

Tensor processing infrastructures such as deep learning frameworks and specialized hardware accelerators have revolutionized how computationally intensive code from domains such as deep learning and image processing is executed and optimized. These infrastructures provide powerful and expressive abstractions while ensuring high performance. However, to utilize them, code must be written specifically using the APIs / ISAs of such software frameworks or hardware accelerators. Importantly, given the fast pace of innovation in these domains, code written today quickly becomes legacy as new frameworks and accelerators are developed, and migrating such legacy code manually is a considerable effort. To enable developers in leveraging such DSLs while preserving their current programming paradigm, we introduce Tenspiler, a verified lifting-based compiler that uses program synthesis to translate sequential programs written in general-purpose programming languages (e.g., C++ or Python code) into tensor operations. Central to Tenspiler is our carefully crafted yet simple intermediate language, named TensIR, that expresses tensor operations. TensIR enables efficient lifting, verification, and code generation. Currently, Tenspiler already supports $ extbf{six}$ DSLs, spanning a broad spectrum of software and hardware environments. Furthermore, we show that new backends can be easily supported by Tenspiler by adding simple pattern-matching rules for TensIR. Using 10 real-world code benchmark suites, our experimental evaluation shows that by translating code to be executed on $ extbf{6}$ different software frameworks and hardware devices, Tenspiler offers on average 105$ imes$ kernel and 9.65$ imes$ end-to-end execution time improvement over the fully-optimized sequential implementation of the same benchmarks.

研究の動機と目的

• 一般コードからドメイン特化言語（DSL）ターゲットへの自動変換を可能にすることで、テンソルコンピューティングにおけるレガシーコードの移植性の課題に対処する。
• パターンマッチングコンパイラや LLM に基づくコード生成の限界を克服し、破綻しやすく、誤りが入りやすく、正しさの保証がない問題を解消する。
• 効率的な合成と検証を可能にするために、テンソル演算を抽象化する統一的かつ拡張可能な中間表現（Tensir）を設計する。
• Tensir に単純なパターンマッチングルールを追加することで、新しいバックエンドへの対応をシームレスに可能にし、新しいフレームワークやアクセラレータへの移植作業を軽減する。

提案手法

• 入力プログラムを同等の Tensir プログラムにリフトすることで、コード変換タスクを構文誘導型プログラム合成（SGS）問題として定式化する。
• 定理証明器を用いて、入力プログラムと合成された Tensir プログラムの意味的同等性を検証し、正しさを保証する。
• 一般的なベクトルおよび行列演算を捉えることのできる最小限のテンソル代数ベースの IR として Tensir を設計し、効率的な合成と検証を可能にする。
• 検索空間の縮小とスケーラビリティの向上を図るため、演算子制限やプログラム状態制限などの最適化技術を適用する。
• ユーザーが提供するコードジェネレータを用いて、検証済みの Tensir プログラムからターゲットコードを生成し、MLX、Gemmini、Gaudi など複数のバックエンドをサポートする。
• 式木とベクタライゼーションヒューリスティクスを活用して、高次元テンソルにおける効率的でベクタライズされた演算の合成をガイドする。

実験結果

リサーチクエスチョン

• RQ1 検証済みのリフトベースのアプローチは、多様なテンソルDSLおよびハードウェアバックエンドへの一般コードの正しくかつ効率的なコンパイルを可能にするか？
• RQ2 統一的中間表現（Tensir）は、異種のテンソル処理ターゲットにわたり、スケーラブルな合成と検証を可能にするか？
• RQ3 Tenspiler の合成と検証パイプラインは、1D、2D、3D テンソル演算を含む、現実世界の複雑なベンチマークにどの程度スケーラブルに拡張できるか？
• RQ4 Tenspiler のパフォーマンスは、完全に最適化された順次実装と比較して、カーネルおよびエンドツーエンド実行時間の点でどの程度優れているか？
• RQ5 Tenspiler は、最小限の工学的作業で新しいバックエンドを容易に拡張できるか？

主な発見

• Tenspiler は、69 個の実世界ベンチマークにおいて、完全に最適化された順次実装と比較して、カーネル実行時間で平均 105× の高速化、エンドツーエンド実行時間で 9.65× の高速化を達成した。
• 合成プロセスは効果的にスケーリングされ、深さ 5 の人工的 3D テンソルベンチマークで 184 秒（1 時間以内のタイムアウト内）で合成が完了した。
• Tenspiler は、MLX、Gemmini、Gaudi を含む 6 つの異なるバックエンドにコードを変換することができ、オープンソースおよび商用フレームワークとの広範な互換性を示した。
• 新規リリースの MLX フレームワーク向けに Tenspiler を実装するのに 200 行未満のコードで完了した。これは、Tensir IR の拡張性の高さを示している。
• LLM が生成したコードは、同じベンチマークで一貫して誤りで機能せず、未熟な DSL に対して検証済みコードを生成する際の LLM の限界を強調している。
• Tensir の使用により、効率的な検証とコード生成が可能となり、演算子制限やプログラム状態制限などの最適化が合成の検索空間を顕著に縮小した。