[論文レビュー] Tensor Processing Primitives: A Programming Abstraction for Efficiency and Portability in Deep Learning & HPC Workloads
本稿では、深層学習およびHPCワークロードの高性能でポータブルな実装を可能にする、最小限の2次元テンソル演算を抽象化したポータブルで効率的なプログラミング抽象化であるテンソル処理プリミティブ（TPP）を紹介する。アルゴリズムの記述とプラットフォーム固有の最適化を分離することで、TPPは開発者やコンパイラが複雑な演算子をポータブルに表現しつつ、複数のアーキテクチャで優れたパフォーマンスを達成できる。エンドツーエンドのDLワークロードおよびスタンドアロンカーネルにおいて、最先端の実装を上回る性能を発揮する。
During the past decade, novel Deep Learning (DL) algorithms, workloads and hardware have been developed to tackle a wide range of problems. Despite the advances in workload and hardware ecosystems, the programming methodology of DL systems is stagnant. DL workloads leverage either highly-optimized, yet platform-specific and inflexible kernels from DL libraries, or in the case of novel operators, reference implementations are built via DL framework primitives with underwhelming performance. This work introduces the Tensor Processing Primitives (TPP), a programming abstraction striving for efficient, portable implementation of DL workloads with high-productivity. TPPs define a compact, yet versatile set of 2D-tensor operators (or a virtual Tensor ISA), which subsequently can be utilized as building-blocks to construct complex operators on high-dimensional tensors. The TPP specification is platform-agnostic, thus code expressed via TPPs is portable, whereas the TPP implementation is highly-optimized and platform-specific. We demonstrate the efficacy and viability of our approach using standalone kernels and end-to-end DL & HPC workloads expressed entirely via TPPs that outperform state-of-the-art implementations on multiple platforms.
研究の動機と目的
- 深層学習ワークロードとハードウェアの急速な進展にもかかわらず、DLプログラミング手法の停滞を解消すること。
- モノリシックで非ポータブルなDLライブラリカーネルや、新規演算子に対して性能が劣るリファレンス実装の限界を克服すること。
- 低レベルのハードウェア知識を必要とせず、効率的でプラットフォーム最適化された実装を可能にするポータブルで生産性の高い抽象化を提供すること。
- 開発者およびコンパイラの両者において、アルゴリズムの記述と低レベルコード生成の間に明確な責任の分離を実現すること。
- 標準化された最小限の演算子セットを通じて、フレームワーク、コンパイラ、ハードウェアプラットフォーム間での再利用と最適化の共有を促進すること。
提案手法
- テンソルの仮想命令セットアーキテクチャ（ISA）として、コンactで精度に配慮した2次元テンソルプリミティブ（TPP）のセットを定義する。
- TPPをプラットフォームに依存しない仕様で設計するとともに、実装段階で高い最適化を実現し、パフォーマンスを最大化するポータブルコードを可能にする。
- TPPを基本的なモジュールとして用い、DLおよびHPCワークロードにおける複雑な高次元演算子を表現する。
- DLフレームワーク、ベンダーライブラリ（例：oneDNN）、テンソルコンパイラ（例：MLIRディアレクト）のソフトウェアスタックにTPPを低下レイヤーとして統合する。
- ターゲットアーキテクチャ向けにインプリント命令（例：AVX512、SVE、NEON）を用いたTPPバックエンドを実装し、低レベルのパフォーマンスを達成する。
- 既存のコンパイラインfra（例：MLIR）を活用し、TPPベースのカーネルのコード生成および最適化を可能にする。
実験結果
リサーチクエスチョン
- RQ1最小限のポータブルな2次元テンソル演算セットが、複雑で高性能なDLおよびHPCワークロードの実装に有効な抽象化として機能できるか？
- RQ2アルゴリズムの記述（TPP経由）とプラットフォーム固有の最適化を分離することで、従来のDLライブラリカーネルや汎用フレームワークプリミティブと比較して、パフォーマンスとポータビリティがどのように向上するか？
- RQ3TPPが、多様なハードウェアプラットフォームにおける実世界のDLおよびHPCワークロードにおいて、最先端の実装をどの程度上回るか？
- RQ4MLIR、DLフレームワーク、ベンダー最適化ライブラリといった既存のソフトウェアスタックにTPPを効果的に統合できるか？
- RQ5仮想テンソルISA（TPP）の使用が、コンパイラやシステム間での最適化の共有を促進し、重複を削減する仕組みをどのように実現するか？
主な発見
- CNN、拡張畳み込み、DLRM、BERT、GNNなどのエンドツーエンドDLワークロードをTPPのみで実装した場合、複数のプラットフォームで最先端の実装を上回る性能を発揮した。
- TPPを用いたスタンドアロンカーネル実装は、最適化されたベンダーライブラリやリファレンス実装を上回る優れたパフォーマンスを達成した。
- TPP抽象化により、低レベルのパフォーマンスを維持しつつ、高い生産性とポータビリティを実現した。これは、プラットフォーム固有のバックエンド最適化のおかげである。
- TPPはMLIRとしてネイティブに統合可能であり、複数のコンパイラスタックで効率的なコード生成と最適化を可能にした。
- データ／ML科学者が手動で低レベル最適化を行う必要が減り、新規演算子の開発とデプロイのスピードアップが実現した。
- x86、AArch64、ARMv8 SVE、および将来のGPU／POWER対応を含む多様なISAのサポートは、TPP抽象化の拡張性とポータビリティを示している。
より良い研究を、今すぐ始めましょう
論文設計から論文執筆まで、研究時間を劇的に削減しましょう。
クレジットカード登録不要
このレビューはAIが作成し、人間の編集者が確認しました。