[論文レビュー] Stripe: Tensor Compilation via the Nested Polyhedral Model

Tim Zerrell, Jeremy Bruestle|arXiv (Cornell University)|Mar 14, 2019

Parallel Computing and Optimization Techniques参考文献 20被引用数 25

ひとこと要約

この論文では、ネストド・ポリヘドラル・モデルに基づく、機械学習コンパイラー向けのドメイン特化中間表現（IR）であるStripeを紹介する。このIRは、多様なハードウェア上で自動的かつ高パフォーマンスなカーネル生成を可能にする。テンソル計算を階層的でポリヘドラルな反復空間としてモデル化することで、アルゴリズム設計、最適化パス、ハードウェアターゲットを分離し、モジュラーで再利用可能な最適化と、現代のMLアクセラレータ向けにスケーラブルなコード生成を実現する。

ABSTRACT

Hardware architectures and machine learning (ML) libraries evolve rapidly. Traditional compilers often fail to generate high-performance code across the spectrum of new hardware offerings. To mitigate, engineers develop hand-tuned kernels for each ML library update and hardware upgrade. Unfortunately, this approach requires excessive engineering effort to scale or maintain with any degree of state-of-the-art performance. Here we present a Nested Polyhedral Model for representing highly parallelizable computations with limited dependencies between iterations. This model provides an underlying framework for an intermediate representation (IR) called Stripe, amenable to standard compiler techniques while naturally modeling key aspects of modern ML computing. Stripe represents parallelism, efficient memory layout, and multiple compute units at a level of abstraction amenable to automatic optimization. We describe how Stripe enables a compiler for ML in the style of LLVM that allows independent development of algorithms, optimizations, and hardware accelerators. We also discuss the design exploration advantages of Stripe over kernel libraries and schedule-based or schedule-space-based code generation.

研究の動機と目的

• 機械学習アルゴリズムとハードウェアアクセラレータの急速な進化に対応するための、手動で最適化されたカーネルの組み合わせ的爆発を解消すること。
• アルゴリズム開発、最適化パス、ハードウェア固有のコード生成を分離することで、保守負荷を低減し、拡張性を向上させること。
• 複雑な並列性とメモリアクセスパターンを構造的かつ拡張可能なIRを用いて形式的にモデル化することで、MLワークロードに対する自動的かつ高パフォーマンスなコード生成を可能にすること。
• 最適化パスが異なる演算やハードウェアプラットフォーム間で再利用可能で合成可能である、LLVMに類似したコンパイラー・スタックの基盤を提供すること。
• 一般用途コンパイラーおよび既存のML特化フレームワークの限界を克服し、数学的に根拠のある階層的表現を導入することで、テンソル計算を形式的に表現すること。

提案手法

• 高度に並列化可能なテンソル計算を、限られた反復間依存性を持つ形式的枠組みとして表現するため、ネストド・ポリヘドラル・モデルを提案する。
• ポリヘドラル制約とネスティング構造を用いて、反復空間、メモリレイアウト、コンピューティングユニットの抽象化を捉えるドメイン特化IRとしてStripeを設計する。
• 階層的タイリングとスケジューリングをサポートするようにIRを構造化し、複数のメモリレベルおよびコンピューティングユニットを横断した最適化を可能にする。
• PlaidMLコンパイラー・スティックにStripeを統合し、Tile（高レベルIR）から低下された高レベルテンソル演算を受信した後、ハードウェアに依存しない最適化パスを適用する。
• 標準的なコンパイラー技術（例：タイリング、ベクタライゼーション、メモリレイアウト変換）をStripe IRに適用し、最適化をターゲットハードウェアに依存せずに独立して適用可能にする。
• ネストド・ポリヘドラル・モデルの形式的構造を活用することで、新しい最適化パスをStripe IR上で構築可能にし、拡張性を実現する。

実験結果

リサーチクエスチョン

• RQ1 複雑で階層的な並列性とメモリアクセスパターンを有する現代のMLワークロードをモデル化しつつ、自動最適化に適したコンパイラーIRをどのように設計できるか？
• RQ2 畳み込みのような演算で一般的に見られる非直交型の反復空間や境界条件を、ネストド・ポリヘドラル・モデルが効果的に表現できるか？
• RQ3 コンパイラー・スタックにおいて、アルゴリズム、最適化、ハードウェアの分離をどの程度達成できるか、性能を損なわずに行えるか？
• RQ4 表現力、モularity、パフォーマンスのポータビリティの観点から、Halide や TVM や MLIR といった既存のアプローチと比較して、Stripe IR はどの程度優れているか？
• RQ5 各演算とアーキテクチャの組み合わせに対して手動で最適化されたカーネルを必要とせずに、多様なハードウェアプラットフォームで自動的かつ高パフォーマンスなコード生成を可能にできるか？

主な発見

• ネストド・ポリヘドラル・モデルは、ホーリー領域や境界制約をポリヘドラル枠組み内で形式的にモデル化することで、非ポイントワイズ畳み込みに見られるような複雑で非直交型の反復空間を効果的に捉えることに成功した。
• Stripeのネストド・ブロック構造は、マルチレベルタイリングを自然にサポートし、ワークロードを階層的メモリシステムおよび複数のコンピューティングユニットに効率的にマッピング可能である。
• IRは、異なるML演算やハードウェアターゲット間で最適化パスを独立して開発・再利用可能にし、手動最適化カーネルの必要性を低減する。
• GPUターゲットにおける初期ベンチマークにより、Stripeは手動最適化実装と競合するパフォーマンスでMLカーネルの自動コード生成を可能にした。
• PlaidMLのコンパイルパイプラインへのStripe統合により、ハードウェアに依存しないコンパイルが可能になり、再利用可能で合成可能な最適化パスが実現した。
• 著者らは、将来的にStripeをMLIRのディアレクトとして統合することで、相互運用性が向上し、コンパイラーエコシステム全体にわたる最適化パスの共有が促進されると予想している。