AIコンパイラ原理4

感謝 up 主 ZOMI ちゃん：https://space.bilibili.com/517221395

なぜ AI コンパイラが必要なのか#

直面している問題#

チャレンジカテゴリ	説明	ケース
演算子の課題	新しい演算子がますます提案され、演算子ライブラリの開発、保守、最適化、テストの作業量が指数的に増加しています。	1. ハードウェアは新しい演算子を実装するだけでなく、ハードウェアに基づいて特性の最適化とテストを行い、ハードウェアの性能を最大限に引き出す必要があります。例えば、畳み込み演算については、GEMM 行列乗算に変換する必要があります。新しく提案された Swish 演算子については、ハードウェアに Swish に対応する実装を追加する必要があります。 2. ハードウェアベンダーは最適化ライブラリをリリースしますが、類似の最適化ライブラリを開発する必要があり、これにより多くの演算子の最適化やラッピング作業が増え、ライブラリに過度に依存し、専用ハードウェアチップの能力を効果的に活用できなくなります。
最適化の課題	専用加速チップの爆発により、性能の移植性が必須となっています。	1. ほとんどの NPU は ASIC を使用しており、神経ネットワークシーンにおいて計算、ストレージ、データ転送に特別な命令最適化を行い、AI 関連計算の性能を向上させています。 2. 異なるベンダーが提供する XPU の ISA は多様で、GCC、LLVM などのコンパイラツールチェーンが不足しているため、CPU および GPU 向けの既存の最適化演算子ライブラリや言語向けの最適化パスを NPU に短期間で移植することが困難です。

従来のコンパイラと AI コンパイラ#

従来のコンパイラと AI コンパイラの主な違いは 2 つあります：

IR の違い：AI コンパイラの IR は従来のコンパイラの IR が抽象化した概念や意味とは異なります。
1. AI コンパイラは一般に high-level IR を持ち、深層学習モデル内の演算を抽象的に記述するために使用されます。例えば、畳み込み、行列乗算など、場合によっては Transformer を含むグラフ構造を持つこともあります。
2. 従来のコンパイラは比較的 low-level IR を使用し、基本的な命令演算を記述します。例えば、ロード、ストアなど。high-level IR があることで、AI コンパイラは深層学習モデルの DSL をより便利に記述できます。
最適化戦略：AI コンパイラは AI 分野に特化しており、最適化時により多くのドメイン特有の知識を取り入れ、より high-level で、より攻撃的な最適化手法を行います。例えば：
1. AI コンパイラは high-level IR で演算子の融合を実行し、従来のコンパイラがループ融合を実行する際には、しばしばより保守的です。欠点は、デバッグ実行情報の追跡が難しくなる可能性があることです；
2. AI コンパイラは計算精度を低下させることができます。例えば、int8、fp16、bf16 など、深層学習は計算精度にそれほど敏感ではありません。しかし、従来のコンパイラは一般に変数の型や精度を変更する最適化を実行しません。

AI コンパイラの発展段階#

推論シーン：AI フレームワークでトレーニングされたモデルファイルを入力し、異なるハードウェアで効率的に実行できるプログラムを出力します；
トレーニングシーン：高級言語で表現された神経ネットワークコードを入力し、異なるハードウェアで効率的に実行できるプログラムを出力します；

トレーニングシーンとは何ですか？推論シーンとは何ですか？
AI コンパイラを扱うにはなぜアルゴリズムを理解する必要があるのですか？
AI 演算子を扱うにはなぜコンパイラを理解する必要があるのですか？

AI コンパイラとは#

Python を主とした動的インタプリタ言語のフロントエンド
複数の IR 設計、グラフコンパイル、演算子コンパイル、コード生成を含む
神経ネットワーク、深層学習に特化した最適化
DSA チップアーキテクチャのサポート

AI コンパイラの発展段階#

ステージ I：素朴な AI コンパイラ#

TensorFlow の初期バージョンで、神経ネットワークのプログラミングモデルに基づき、主にグラフと演算子の 2 層の抽象化が行われました。

グラフ層：宣言的なプログラミング方式で静的グラフ方式で実行し、実行前にハードウェアに依存しない最適化とハードウェアに依存する最適化を行います。ハードウェアに依存しない最適化には、式の簡略化、定数折りたたみ、自動微分などが含まれます。ハードウェアに依存する最適化には、演算子の融合、メモリ割り当てなどが含まれます。
演算子層：通常、手書きのカーネルを使用します。例えば、NVIDIA GPU 上で CUDA カーネルに基づいて多数の.cu演算子を実装したり、CuDNN 演算子最適化ライブラリに依存したりします。

表現において：

静的グラフの表現は Python のネイティブではなく、開発者は主にフレームワークが提供する Python API を通じてグラフを表示し、使いやすさが悪いです；

性能において：

DSA 専用加速チップの出現が性能上の課題を悪化させました；
演算子層が提供する演算子の粒度と境界が事前に決定されると、ハードウェアの性能を十分に発揮できません；
ハードウェアベンダーが提供する演算子最適化ライブラリも必ずしも最適ではありません
- 1）モデルと形状が確定した場合、より優れた演算子実装がある可能性があります；
- 2）SIMT および SIMD アーキテクチャにおいて、スケジューリング、ティリングには大きな余地があります。

ステージ II 専用の AI コンパイラ#

表現において：

PyTorch の柔軟な表現 API 方式が AI フレームワークの参考基準となり、グラフ層の神経ネットワークコンパイラは主に PyTorch のような表現をグラフ層の IR に変換して最適化することを考慮しています。
PyTorch のような Python ネイティブ表現、静的化変換；
AI 専用コンパイラアーキテクチャがグラフと演算の境界を開いて融合最適化を行います；

性能において：

計算グラフと演算子の境界を開き、再構成最適化を行い、チップの計算能力を発揮します。計算グラフ層はサブグラフ内の演算子を小さな演算子に開き、小さな演算子で構成されたサブグラフに対してコンパイル最適化を行います。これにはバッファ融合、水平融合などが含まれ、重要なのは大きな演算子をどのように開き、小さな演算子をどのように再融合するかです。

表現の分離：計算グラフ層と演算子層は依然として分かれており、アルゴリズムエンジニアは主にグラフ層の表現に注目し、演算子の表現と実装は主にフレームワーク開発者とチップベンダーが提供します。
機能の一般化：柔軟な表現における動的・静的グラフの変換、動的形状、疎計算、分散並列最適化などの複雑な要求を満たすのが難しいです。
効率と性能のバランス：演算子の実装においてスケジュール、ティリング、コード生成に自動化手段が不足しており、ハードルが高く、開発者は演算子の計算ロジックを理解するだけでなく、ハードウェアアーキテクチャにも精通する必要があります。

ステージ III 汎用 AI コンパイラ#

グラフと演算の統一表現を実現し、融合最適化を行います
演算子の実装において自動スケジュール、ティリング、コード生成を行い、開発のハードルを下げます
より一般化された最適化能力を実現し、動的・静的の統一、動的形状、疎性、高次微分、自動並列などを実現します
コンパイラ、ランタイムを含み、異種計算、エッジからデータセンターまでモジュール化された表現と組み合わせを行い、使いやすさに焦点を当てます

AI コンパイラの汎用アーキテクチャ#

IR 中間表現#

コンパイラは主にフロントエンドとバックエンドに分かれており、それぞれハードウェアに依存しない処理とハードウェアに依存する処理を行います。各部分にはそれぞれの IR があり、各部分も最適化を行います：

High-level IR：計算グラフを表現するために使用され、従来のコンパイラでは深層学習モデル内の複雑な演算を表現するのが難しいという問題を解決するために登場しました。より効率的な最適化を実現するために新たに設計された IR です。
Low-level IR：より細かい粒度でモデルを表現できるため、ハードウェアに対して最適化を行うことができます。文中ではこれを 3 つのカテゴリに分けています。

フロントエンド最適化#

計算グラフを構築した後、フロントエンドはグラフレベルの最適化を適用します。グラフは計算の全体像を提供するため、グラフレベルで多くの最適化を発見し実行するのが容易です。フロントエンド最適化はハードウェアに依存しないため、計算グラフの最適化をさまざまなバックエンドターゲットに適用できます。フロントエンド最適化は 3 つのカテゴリに分けられます：

ノードレベルの最適化、例えばゼロ次元テンソルの排除、ノップの排除
ブロックレベルの最適化、例えば代数の簡略化、定数の折りたたみ、演算子の融合
データフローレベルの最適化、例えば共通部分式の排除、DCE

バックエンド最適化#

特定ハードウェアの最適化
- 特定のハードウェアアーキテクチャに対して高性能なコードを取得することを目指します。1）低レベル IR を LLVM IR に変換し、LLVM 基盤を利用して最適化された CPU/GPU コードを生成します。2）ドメイン知識を使用して最適化をカスタマイズし、ターゲットハードウェアをより効果的に活用します。
自動調整
- 特定のハードウェア最適化においてパラメータ調整に使用される探索空間が巨大であるため、最適なパラメータ設定を決定するために自動調整を利用する必要があります。1）Halide/TVM はスケジューリングと計算表現を分離し、自動調整を使用して最適な構成を導き出します。2）多面体モデルを適用してパラメータ調整を行います。
最適化カーネルライブラリ
- ベンダー特定の最適化カーネルライブラリは、さまざまなハードウェア上で DL トレーニングと推論を加速するために広く使用されます。特定の最適化原語が計算要求を満たす場合、最適化されたカーネルライブラリを使用することで性能が大幅に向上しますが、そうでない場合はさらなる最適化の制約を受ける可能性があります。

AI コンパイラの課題と考察#

XLA：機械学習コンパイラの最適化#

XLA（加速線形代数）は特定のドメイン向けの線形代数コンパイラで、TensorFlow モデルの実行速度を加速し、ソースコードを変更することなく実行できます。

XLA：グラフ層はサブグラフ内の演算子を小さな演算子に開き、小さな演算子で構成されたサブグラフに対してコンパイル最適化を行います。これにはバッファ融合、水平融合などが含まれ、重要なのは大きな演算子をどのように開き、小さな演算子をどのように再融合し、新しい大きな演算子をどのように生成するかです。全体の設計は主に HLO/LLO/LLVM IR を通じて実現され、すべての Pass ルールは手動で事前に指定されます。

TVM：エンドツーエンドの深層学習コンパイラ#

さまざまなハードウェアバックエンドの計算グラフ層と演算子層の最適化を可能にするために、TVM は既存のフレームワークから DL プログラムの高レベル表現を取得し、複数のハードウェアプラットフォームバックエンドで低レベルの最適化コードを生成します。その目標は、人工調整との競争力を示すことです。

TVM：Relay と TVM の 2 層に分かれており、Relay はグラフ層に焦点を当て、TVM は演算子層に焦点を当て、フロントエンドのサブグラフを最適化します。Relay は演算子間の融合に注目し、TVM は新しい演算子とカーネルの生成に注目します。違いは、TVM がオープンアーキテクチャであり、Relay の目標はさまざまなフロントエンドを接続できることです。TVM は独立して使用できる演算子開発とコンパイルのツールでもあり、演算子実装において Compute（計算ロジックの設計）と Schedule（スケジューリング最適化ロジックの指定）を分離したアプローチを採用しています。

TENSOR COMPREHENSIONS：神経ネットワーク言語コンパイラ#

TensorComprehensions は、JIT システムを構築できる言語で、プログラマーはこの言語を使用して高級プログラミング言語で GPU などの低レベルコードを効率的に実装できます。

TC：演算子計算ロジックの実装は比較的容易ですが、スケジュールの開発は難しく、アルゴリズムロジックを理解するだけでなく、ハードウェアアーキテクチャにも精通する必要があります。また、計算グラフの境界が開かれ、小さな演算子が融合されると、新しい演算子とカーネルが生成されますが、新しい演算子のスケジュールを生成するのは難しいです。従来の方法ではスケジュールテンプレートを定義します。TC は多面体モデルを通じて自動スケジュールを実現することを望んでいます。

nGraph：すべてのフレームワークに対応する深層学習システムコンパイラ#

nGraph の運用は、深層学習フレームワーク内のより複雑な DNN 演算の基盤として機能し、推論計算とトレーニングおよび推論計算の間で理想的な効率のバランスを確保します。

課題#

動的形状と動的計算グラフ：
- 現状：主流の AI コンパイラは主に特定の静的形状入力に対してコンパイル最適化を行い、制御フローセマンティクスを含む動的計算グラフのサポートは限られています。
- 問題：AI アプリケーションシーンには動的計算グラフが必要なタスク要求が多数存在します。フロントエンドは計算グラフを静的計算グラフに書き換えたり、コンパイルに適した部分のサブグラフを展開して最適化を行うことができますが、これらの方法ではすべての問題を解決できません。
- 例：特定の AI タスク（例えば、ピラミッド構造の検出モデル）は人工的に書き換えて静的化することができず、コンパイラはこれらの状況で効果的に最適化するのが難しいです。

Python コンパイルの静的化：

方法	説明
Python JIT 仮想機	PyPy や CPython のように、Python の解釈実行に基づいて JIT コンパイルの加速能力を追加しようとしますが、互換性に問題があります。
修飾子方式	torch.jit.script、torch.jit.trace、torch.fx、LazyTensor、TorchDynamo などのように、これらの手段は広く使用されていますが、完璧なソリューションはまだありません。

Python JIT 仮想機

**修飾子方式**

- **AIコンパイラにおけるPython静的化の課題**：
    - **型推論**：Pythonの動的型からコンパイラIRの静的型への変換。
    - **制御フロー表現**：if、else、while、forなどの制御文の表現。
    - **柔軟な構文とデータ型変換**：Slice、Dict、Indexなどの操作の処理。
    - **JITコンパイル性能**：Tracing BasedでもAST Transformでも、追加のコンパイルオーバーヘッドが必要です。

3. ハードウェア性能の発揮、特に DSA タイプのチップ：

現状：DSA（専用ハードウェアアーキテクチャ）は AI トレーニングと推論においてますます重要になっています。例えば、CPU の SIMD ユニット、GPU の SIMT アーキテクチャ、Huawei Ascend の Cube コアなど。
課題：
- 性能最適化はグラフと演算の融合に依存：グラフ層と演算子層を独立して最適化する必要があり、チップ性能を十分に発揮できません。グラフと演算の融合最適化が必要です。例えば、サブグラフの分割、サブグラフ内の垂直融合最適化、水平並列最適化など。
- 最適化の複雑さの増加：スカラー、ベクトル、テンソル、加速命令の多層ストレージ構造が関与し、カーネル実装のスケジューリング、ブロッキング、ベクトル化、テンソル化が複雑になります。
解決策：
- グラフと演算の境界を開いて再構成最適化を行い、チップ性能を十分に発揮します。
- 様々な最適化手法を融合します：垂直融合最適化（例えばバッファ融合）と水平並列最適化（例えばデータ並列）。
- 再構成最適化後のサブグラフのカーネルコードを自動生成します。例えば、スケジューリング、ブロッキング、ベクトル化など。

神経ネットワークの特性の処理：自動微分、自動並列など：

タスク	現状	課題
自動並列	現在の大規模モデルのトレーニングはメモリの壁と性能の壁に直面しており、複雑な並列戦略が必要です。 Scale out：データ並列、テンソル並列、パイプライン並列などの多次元混合並列能力。 Scale up：再計算、混合精度、異種並列など。	手動で設定された分割戦略に依存し、ハードルが高く、効率が低いです。半自動並列は一部の効率問題を解決できますが、最適な並列戦略を自動的に見つけるにはコンパイルと凸最適化問題の解決が必要です。
自動微分	制御フロー：動的グラフは Python 側で制御フローを実行し、ループ回数が多いと性能が劣化します；静的グラフの自動微分は論理の接続や計算グラフの展開方式で解決され、ある程度性能問題を解決できますが、方案はまだ改善が必要です。高次微分：ヤコビ行列を通じて高次前向きおよび後向き微分を効果的に展開し、ヘッセ行列を通じて高次微分を迅速に計算します。	制御フロー：動的グラフは複雑な制御フローを処理する際に性能が低下し、特にループ回数が多い場合に顕著です。静的グラフ方案は性能と安定性を向上させるためにさらなる最適化が必要です。高次微分：高次前向きおよび後向き微分を効果的に展開し計算する方法、計算の正確性と効率を確保する必要があります。開発者は高次微分グラフを柔軟に制御する必要があります。

使いやすさと性能の両立

面	説明
AI フレームワークとの境界と接続	異なる AI フレームワークは深層学習タスクの抽象的な記述と API インターフェースが異なり、それぞれ特徴があります。すべての演算子を完全にサポートしない場合でも、ユーザーの計算グラフの記述を透明にサポートする方法を考慮する必要があります。
ユーザーの透明性の問題	一部の AI コンパイラは完全自動のコンパイラツールではなく、性能はユーザーが提供する高レベルの抽象実装テンプレート（例えば TVM）に依存します。これにより手動調整の演算子実装の人件費が削減されますが、既存の抽象は革新的なハードウェアアーキテクチャに必要な演算子実装を十分に記述できず、コンパイラアーキテクチャに十分精通していないと二次開発やアーキテクチャの再構築が難しく、ハードルが高く、開発負担が重くなります。
コンパイルオーバーヘッド	AI コンパイラは性能最適化ツールとして、コンパイルオーバーヘッドが性能利益に対して十分に優位である場合にのみ実用的な価値があります。特定のアプリケーションシーンでは、コンパイルオーバーヘッドに対する要求が高く、AI コンパイラを使用することでモデルのデバッグと検証を迅速に完了することが妨げられ、開発と展開の難易度と負担が増加する可能性があります。
性能問題	コンパイラの最適化は本質的に人工最適化手法や探求が難しい最適化手法を一般化と抽象化を通じて、限られたコンパイルオーバーヘッドで手動最適化の人件費を代替するものです。深層学習コンパイラは性能が実際に人工最適化を代替または超えることができる場合にのみ価値を発揮します。
堅牢性	大多数の AI コンパイラは研究段階にあり、製品の成熟度は産業レベルのアプリケーションとは大きなギャップがあります。計算グラフのコンパイルがスムーズに行えるか、計算結果の正確性を確保できるか、エラーの追跡とデバッグの能力を考慮する必要があります。

その他の問題#

グラフと演算を統一して表現し、統一的なコンパイル最適化を形成し、汎用 AI コンパイラを作成できるか？
現在の AI フレームワークでは、グラフ層と演算子層は分かれて表現され、最適化されています。アルゴリズムエンジニアは主にグラフ層の表現に接触し、AI フレームワーク、チップ、カーネル開発エンジニアは主に演算子の表現を担当しています。将来的には IR がグラフと演算の間のギャップを打破し、AI + 科学計算、AI + ビッグデータなどの新しいシーンの推進により、計算グラフ層と演算子層が明確でなくなり、AI のコンパイル最適化を統一できるかどうかが問われます。
完全な自動並列は実現可能か？
自動並列はユーザーが入力した直列ネットワークモデルと提供されたクラスターリソース情報に基づいて自動的に分散トレーニングを行うことができます。統一された分散計算グラフと統一されたリソースグラフ設計を採用することで、任意の並列とさまざまなハードウェアクラスターリソース上での分散トレーニングをサポートし、全体コストモデルに基づくプランナーを利用してトレーニングタスクに対してハードウェア感知の並列戦略を自動的に選択できます。実際、自動並列は戦略探索の問題であり、有限の探索空間内で次善の答えを見つけることができますが、真の意味での自動並列が実現できるかどうかはさらなる考察と検証が必要です。
AI チップはコンパイラを必要とするか？AI チップは AI コンパイラを必要とするか？
AI チップがコンパイラに依存するかどうかはチップ自体の設計に依存します。柔軟なチップほどコンパイラへの依存度が高くなります。AI チップの設計初期には CISC スタイルでチップ内部で最適化を解決しました。しかし、専用分野の進化に伴い、より柔軟な要求をサポートするために、AI チップ自体はテンソル命令セットと特殊なメモリ構造を保持しつつ、ますます柔軟になります。将来のアーキテクトはチップとシステムの協調設計を行い、自動化もますます専用チップに適用されるようになるでしょう。

未来#

コンパイラの形態：推論とトレーニングを分け、AOT と JIT の 2 つのコンパイル方式が共存します
IR の形態：AI に統一された IR 表現を持つ MLIR のようなものが必要です
自動並列：機械間、ノード間の自動並列のコンパイル最適化能力を提供します
自動微分：高次微分の計算方法を提供し、グラフの操作を容易にします
カーネル自動生成：開発のハードルを下げ、高効率で高汎用の演算子を迅速に実現します