[論文レビュー] The CoRa Tensor Compiler: Compilation for Ragged Tensors with Minimal Padding

Pratik Fegade, Tianqi Chen+2 more|arXiv (Cornell University)|Oct 19, 2021

Parallel Computing and Optimization Techniques被引用数 8

ひとこと要約

CoRa は、最小限のパディングで不規則なデータ形状を効率的に処理できるテンソルコンパイラであり、CPU および GPU 上での最適化コード生成により、不規則なテンソル形状の演算を効率的に行える。NVIDIA V100 GPU では PyTorch より 1.6× のジオメトリック平均の高速化を達成し、64コアの ARM CPU では TensorFlow より 1.37× の高速化を達成した。

ABSTRACT

There is often variation in the shape and size of input data used for deep learning. In many cases, such data can be represented using tensors with non-uniform shapes, or ragged tensors. Due to limited and non-portable support for efficient execution on ragged tensors, current deep learning frameworks generally use techniques such as padding and masking to make the data shapes uniform and then offload the computations to optimized kernels for dense tensor algebra. Such techniques can, however, lead to a lot of wasted computation and therefore, a loss in performance. This paper presents CoRa, a tensor compiler that allows users to easily generate efficient code for ragged tensor operators targeting a wide range of CPUs and GPUs. Evaluating CoRa on a variety of operators on ragged tensors as well as on an encoder layer of the transformer model, we find that CoRa (i)performs competitively with hand-optimized implementations of the operators and the transformer encoder and (ii) achieves, over PyTorch, a 1.6X geomean speedup for the encoder on an Nvidia GPU and a 1.86X geomean speedup for the multi-head attention module used in transformers on an ARM CPU.

研究の動機と目的

• 可変サイズのデータを処理するディープラーニングフレームワークにおけるパディングとマスキングに起因するパフォーマンスボトルネックを解消すること。
• CPU や GPU といった多様なハードウェア基盤上で、不規則テンソル演算を効率的かつポータブルにコンパイルすること。
• 不規則なテンソル演算のための高パフォーマンスコード生成を自動化することで、手作業でのカーネル最適化の必要性を低減すること。
• 従来の密行列およびスパーステンソルコンパイラでは、不規則テンソル演算を効果的に表現・最適化できないという限界を克服すること。
• CPU および GPU ベースの両方をサポートする統合的かつ拡張可能なコンパイラスタックを提供し、形状の不規則性に起因するパフォーマンス損失を最小限に抑えること。

提案手法

• 可変サイズのテンソルスライスと不規則なループ境界を表現できるドメイン固有言語（DSL）を設計し、不規則テンソル演算を記述する。
• 高レベルの演算記述を最適化された LLVM や CUDA コードに変換するコード生成パイプラインを導入し、形状に配慮したスケジューリングによりパディングを回避する。
• タイリング、統合、並列化などのループ変換技術を活用し、可変なループ境界を扱いながらもデータ局所性とインstruction-level parallelism（命令レベル並列性）を維持する。
• Z3 を用いた制約解決と LLVM を用いた低レベルコード生成を統合し、メモリアクセスパターンやカーネル起動設定の自動最適化を可能にする。
• CPU および GPU ベースの両方をサポートするため、ハードウェアに適した最適化（例：レジスタブロッキング、ワープレベルプリミティブ）を備えたターゲット固有のコード生成を実施する。
• モジュラーなアーキテクチャを採用し、演算子の定義とコード生成を分離することで、新しい演算子やハードウェアプラットフォームへの容易な拡張を可能にする。

実験結果

リサーチクエスチョン

• RQ1 パディングやマスキングに依存せずに、不規則テンソル演算のための高パフォーマンスコードをテンソルコンパイラが生成できるか？
• RQ2 CoRa が自動生成したコードのパフォーマンスは、多様なハードウェアプラットフォームにおける手作業最適化実装と比べてどの程度の水準にあるか？
• RQ3 トランスフォーマーモデルにおいて、パディングベースのアプローチと比較して、CoRa は計算の無駄をどの程度削減できるか？
• RQ4 CoRa は、パフォーマンスとポータビリティ、実装コストの低さを両立させながら、CPU および GPU の両方のワークロードで競争力のあるパフォーマンスを達成できるか？
• RQ5 CoRa は、マルチヘッドアテンションやトランスフォーマーエンコーダーレイヤーのような複雑な実世界の演算子を、どの程度効果的に最適化できるか？

主な発見

• NVIDIA V100 GPU 上でトランスフォーマー・エンコーダーレイヤーに対して、PyTorch より 1.6× のジオメトリック平均の高速化を達成し、手作業最適化カーネルと同等の競争力を持つことが示された。
• 64コアの ARM Neoverse N1 CPU 上でマルチヘッドアテンションモジュールに対して、TensorFlow より 1.37× のジオメトリック平均の高速化を達成し、ベースラインを顕著に上回った。
• CoRa が生成したカーネルのパフォーマンスは、手作業最適化実装と同等であり、自動コード生成が熟練した最適化と同等の水準に達していることを示している。
• パディングを回避することで無駄な計算を削減し、特にバッチサイズが大きい場合に顕著になるパディングのオーバーヘッドが低減されている。トランスフォーマー・エンコーダーの FLOP 分析からもその効果が裏付けられている。
• GPU などの高度に並列化されたアーキテクチャですら、パフォーマンスを損なわず、不規則なループ構造を効果的に処理できることを実証した。
• 8つの NLP データセットを用いた評価により、さまざまなシーケンス長やデータ分布に対しても CoRa のパフォーマンスが安定しており、一般化能力が裏付けられた。