PolyBlocks: A Compiler Infrastructure for AI Chips and Programming Frameworks

本論文は、MLIR ベースのモジュール型コンパイラ基盤「PolyBlocks」の設計と実装を提案し、その自動コード生成能力が PyTorch や JAX 向けに NVIDIA GPU で Torch Inductor や XLA と同等、あるいはそれ以上の性能を発揮し、個別演算においてもベンダー最適化ライブラリや手書きカーネルと競合する結果を示したものである。

要約

PolyBlocks：AI 用「万能調理ロボット」の設計図

この論文は、**「PolyBlocks（ポリブロック）」という新しいシステムの紹介です。これを一言で言うと、「AI の計算を、どんな新しいハードウェア（チップ）でも、自動で最高性能にしてくれる『魔法のコンパイラ』」**です。

少し難しい話になりがちですが、料理とレストランの例えを使って、わかりやすく解説します。

---

1. 今までの問題：「レシピ」と「厨房」の壁

AI を動かすには、大きく分けて 3 つのレベルがあります。

1. 高級レストラン（高レベル：PyTorch, JAX など）

   • 料理人（開発者）は「この材料で、あの料理を作って」という大まかなレシピを書くだけです。
   • しかし、実際の料理（計算）は、プロの料理人が手作業で調理するか、あらかじめ作られた**「冷凍食品（ベンダーライブラリ）」**を使う必要があります。
   • 問題点： 冷凍食品は特定の厨房（NVIDIA の GPU など）にしか使えません。新しい厨房を作ろうとすると、また冷凍食品を作る必要があり、非常に手間がかかります。また、レシピを一つ一つ調理するより、**「一気通貫で調理する（融合）」**方が速いのに、冷凍食品を使うとそれができません。

2. 中級厨房（中レベル：Triton など）

   • 料理人が「この食材はここで切る、この鍋で炒める」という詳細な指示を書く必要があります。
   • 性能は良いですが、毎回指示を書くのは大変で、ミスも起きやすいです。

3. 低レベル厨房（低レベル：CUDA など）

   • 料理人が「包丁を 3 回振る、火加減を 5 にする」という極端に細かい指示を書く必要があります。
   • 最高性能が出ますが、プロの職人（ハイレベルなエンジニア）しかできません。

今の課題：
新しい AI チップ（新しい厨房）を作りたいとき、既存のシステムは「冷凍食品（ライブラリ）」に頼りすぎていて、新しい厨房には対応できません。また、冷凍食品を使わないで「ゼロから調理（コード生成）」しようとすると、プロの料理人（手書きの最適化コード）に勝てないというジレンマがありました。

---

2. PolyBlocks の登場：「万能調理ロボット」

PolyBlocks は、この問題を解決する**「完全自動化の調理ロボット」**です。

• 高レベルのレシピを受け取る：
  PyTorch や JAX などの高級レシピ（高レベル言語）をそのまま受け取ります。
• 自動で最適化して調理：
  「冷凍食品」を使わず、ゼロから最高の料理（コード）をその場で作り出します。
• どんな厨房でも動ける：
  厨房（ハードウェア）が変わっても、ロボットの「調理手順（変換パス）」を少し直すだけで、新しい厨房でも最高性能を出せます。

具体的な魔法（技術的な仕組み）

PolyBlocks がやっていることを料理に例えると以下のようになります。

• 「融合（Fusion）」の魔法：

  • 例え： 「玉ねぎを炒める」→「肉を入れる」→「ソースをかける」という工程を、**「一つの鍋で、一度も食材を移さずに」**調理する技術です。
  • 効果： 食材（データ）を冷蔵庫（メモリ）から出して、鍋（チップ）に戻す回数が激減し、調理が爆速になります。既存のシステムは、この「一度に調理する」のが苦手でしたが、PolyBlocks は得意です。

• 「タイル化（Tiling）」の魔法：

  • 例え： 巨大なピザを一度に焼くのではなく、一口サイズ（タイル）に切って、小さなオーブン（オンチップメモリ）で次々と焼く技術です。
  • 効果： 大きな冷蔵庫から取り出す手間を減らし、小さなオーブン（高速なメモリ）の中で効率よく調理できます。

• 「Attention（注意機構）」の最適化：

  • 例え： 最近の AI（LLM など）は、文章の「どこに注目するか」を計算する「Attention」という工程があります。これは通常、非常に時間がかかります。PolyBlocks は、この工程を**「計算とメモリの移動を完全に重ね合わせ」**、無駄な動きをゼロにします。
  • 結果： 既存の最高性能のライブラリ（FlashAttention など）と同等か、それ以上の速さを出します。

---

3. 実験結果：「冷凍食品」に勝つ

このロボットを実際にテストした結果は驚くべきものでした。

• 既存の「冷凍食品」に匹敵：
  NVIDIA 製の GPU 向けに、世界中のプロが何年もかけて作り込んだ「冷凍食品（CuDNN や CuBLAS などのライブラリ）」を使わずに、PolyBlocks がゼロから作った料理は、それらと同等か、それ以上の速さでした。
• 新しい厨房への対応：
  既存のシステムは、新しいチップが出ると「冷凍食品」を作るまで時間がかかりますが、PolyBlocks は**「調理手順（コンパイラ）」**を少し直すだけで、新しいチップでもすぐに最高性能を出せます。

まとめ

PolyBlocksは、AI 開発者が「新しいチップを作りたい」「既存のライブラリに依存したくない」と悩む時に、**「高レベルのレシピを、自動で、最高性能の低レベルな料理に変換してくれる、再利用可能な魔法のコンパイラ」**です。

これにより、AI の未来は「特定のメーカーの冷凍食品に依存する」時代から、「どんな新しいハードウェアでも、自動で最高性能を発揮できる」時代へと変わっていく可能性があります。

技術概要

PolyBlocks: AI チップおよびプログラミングフレームワーク向けのコンパイラ基盤に関する技術的サマリー

本論文は、AI 向けプログラミングフレームワーク（PyTorch, JAX など）と専用 AI チップを接続するための、モジュール性が高く再利用可能なコンパイラ基盤「PolyBlocks」の設計と実装について述べています。PolyBlocks は MLIR（Multi-Level Intermediate Representation）に基づいており、完全なコード生成アプローチを採用することで、ベンダーライブラリに依存せず、高性能なコードを自動的に生成することを可能にします。

以下に、問題定義、手法、主要な貢献、実験結果、および意義について詳細をまとめます。

1. 背景と問題定義

現在の AI 開発エコシステムには、以下の課題が存在します。

• 実行モードの乖離: PyTorch や JAX などの高レベルフレームワークでは、デフォルトの「Eager 実行（逐次実行）」が主流ですが、これは各演算子ごとに事前に最適化されたベンダーライブラリ（cuDNN, cuBLAS など）を呼び出す形式です。これにより、演算子間のフュージョン（結合）が制限され、メモリアクセスのオーバーヘッドが発生します。
• 既存コンパイラの限界: 現在主流の JIT コンパイラ（Torch Inductor, XLA）は、高レベルの最適化を行いますが、計算集約的な演算子（畳み込み、行列積、アテンション層など）の実装において、依然として手書きのベンダーライブラリや中レベルのフレームワーク（Triton など）に依存しています。
  • 移植性の欠如: ライブラリ依存は、新しいハードウェアや異なる問題サイズへの対応を困難にし、ポータビリティを損ないます。
  • フュージョンの制約: ライブラリ呼び出しを介するため、演算子間の完全なフュージョンが実現できず、グローバルメモリへのアクセス回数が減りません。
• 低レベルプログラミングの難易度: 手動で CUDA や Triton を使って最適化を行うには、深いハードウェア知識と多大な労力が必要であり、スケーラビリティに欠けます。

2. 手法とアーキテクチャ

PolyBlocks は、MLIR インフラストラクチャに基づき、高レベルの抽象度からターゲット固有のイントリンシクスまでを完全にコード生成するアプローチを採用しています。

2.1 設計思想

• 完全自動化: プログラマによるスケジュール指定や自動チューニングに依存せず、コンパイラがコストモデルとヒューリスティクスを用いて最適化を決定します。
• 完全コード生成: 手書きライブラリに依存せず、ターゲット固有のインストルメントまで全てを生成します。これにより、新しいハードウェアへの対応が容易になります。
• 再利用可能なパスパイプライン: 変換を「パス（Pass）」の集合として構成し、MLIR のパスパイプラインとして実装しています。これにより、異なるフレームワークやハードウェア間での基盤の再利用が可能になります。
• アフィン解析の活用: 複雑な多項式解析（Polyhedral）ではなく、軽量なアフィン（Affine）アクセス解析を主に使用し、コンパイル時間の短縮とスケーラビリティを両立させています。

2.2 アーキテクチャ（5 段階のパイプライン）

コンパイラは以下の 5 つのステージに整理されています：

1. S1 (Frontend): 高レベルのテンソル操作からメモリ参照（memref）への低レベル化（バッファセマンティクスの導入）。
2. S2: メモリ参照上の名前付き操作から、アフィンループネス（Affine Loop Nests）への低レベル化。
3. S3 (Mid-level Optimizer): 中核となる最適化ステージ。ターゲット情報に基づき、以下の最適化を実行します。
   • マルチレベルのタイリング（Tiling）
   • 演算子間のフュージョン（Fusion）
   • オンチップスクラッチパッド（共有メモリ）へのデータ移動コード生成
   • 行列演算ユニット（Tensor Core など）へのマッピング
   • ベクトル化
4. S4: アフィン形式から並列化抽象（SCF, GPU dialect など）への変換。
5. S5 (Backend): 最終的なターゲット固有の IR（NVVM, LLVM など）への低レベル化とコード生成。

2.3 主要な最適化技術

• スライシングベースのアフィンフュージョン: 従来のループフュージョンを超え、生産者と消費者の間の依存関係を解析し、必要なデータスライスを抽出してフュージョンを行います。これにより、中間バッファの排除やレジスタ再利用が可能になります。
• アテンション層の完全フュージョン: Transformer モデルの Attention 層（QK^T, Softmax, VSoftmax など）を、DRAM やオンチップメモリへの往復を最小化するように自動的にフュージョンします。特に「reduce-reduce-fusion」と「wmma-fusion」のパスを用いて、大規模な中間結果の生成を回避し、レジスタレベルでの再利用を実現しています。
• 畳み込みのオンザフライ・パッキング: 畳み込みを行列積（GEMM）に変換する際、メモリオーバーヘッドを避けるため、高速メモリ上でタイル単位に動的にデータのパッキング（再構成）を行うコードを生成します。
• 行列ユニットへのマッピング: 畳み込みや行列積を、ハードウェア固有の行列演算ユニット（NVIDIA の WMMA など）に自動的にマッピングするパスを提供します。

3. 主要な貢献

1. MLIR 基盤の完全コード生成コンパイラ: ベンダーライブラリに依存せず、PyTorch や JAX から直接、ハードウェア固有の高性能コードを生成する初の包括的な基盤の提示。
2. 高度な自動最適化: 従来のコンパイラでは困難だった、アテンション層全体のフュージョンや、畳み込みの効率的な GEMM 変換（オンザフライ・パッキング）を自動化。
3. モジュール性と移植性: 新しい AI チップへの対応が容易。S1, S2 はターゲット非依存であり、S3 以降をターゲットに合わせてカスタマイズするだけで、新しいバックエンドを構築可能。
4. パフォーマンス競争力: 手書きの最適化ライブラリ（cuDNN, cuBLAS）や、既存の JIT コンパイラ（Inductor, XLA）と同等、あるいはそれ以上の性能を達成。

4. 実験結果

NVIDIA A100 および A10 GPU 上で、PyTorch と JAX のモデル（Vision, NLP, Transformer）を用いて評価を行いました。

• PyTorch での性能:
  • Batch Size 1 (リアルタイム推定): PolyBlocks は Eager 実行に対して平均 2.15 倍、Torch Inductor に対して 1.4 倍、TensorRT に対して 2.4 倍の高速化を達成。
  • Batch Size 8: Inductor と同等の性能（0.97 倍）を達成。Inductor は大規模バッチ向けに最適化されたライブラリに依存しているため、PolyBlocks がライブラリ依存なしでこれに匹敵することは画期的です。
• JAX での性能:
  • XLA に対して平均 1.15 倍、Eager 実行に対して 2.12 倍の高速化を達成。
• 個々の演算子の性能:
  • 畳み込み: 数百の異なる設定において、cuDNN と同等の性能を達成。約 50 のケースでは cuDNN を 2 倍以上上回りました。
  • 行列積 (Matmul): CuBLAS や手書きの Triton カーネルと同等の性能を達成。
  • アテンション層: 最適化を適用することで、Inductor が使用する手書き Flash Attention カーネルと同等、あるいはそれ以上の性能を自動生成コードで達成しました。
• アブレーション研究:
  • Tensor Core の使用により平均 17 倍の高速化。
  • 演算子間フュージョンにより、最適化済みコードでも 2.87 倍の性能向上。
  • アテンション層の最適化（RR-fusion, wmma-fusion）が、大規模モデルにおいて劇的な性能向上をもたらすことが確認されました。

5. 意義と結論

PolyBlocks は、AI コンパイラ開発のパラダイムシフトを示唆しています。

• ライブラリ依存からの脱却: 手書きライブラリに依存しないことで、新しいハードウェアへの迅速な対応と、柔軟な最適化が可能になります。
• 高レベル抽象度の維持: プログラマは高レベルのコード（PyTorch/JAX）を書くだけで、コンパイラがハードウェアの特性（メモリアーキテクチャ、行列ユニットなど）を自動的に活用するコードを生成します。
• 将来のハードウェアへの適応: 既存のライブラリエコシステムが存在しない将来の AI チップに対しても、PolyBlocks の基盤を活用することで、高性能なコード生成が可能になります。

結論として、PolyBlocks は、高レベルフレームワークから低レベルのハードウェアインストルメントまでを完全にカバーし、ベンダーライブラリに匹敵する、あるいは凌駕する性能を自動生成する、実用的でスケーラブルなコンパイラ基盤として確立されました。